#include #include #include #include #include #include using namespace std; /* This TreeNode definition is from Leetcode. */ struct TreeNode { int val; TreeNode *left; TreeNode *right; TreeNode() : val(0), left(nullptr), right(nullptr) {} TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} }; /* I have added the dfs() method and the V data structure to the class definition. */ class Solution { public: vector > verticalTraversal(TreeNode* root); void dfs(TreeNode *root, int r, int c); map > > V; }; /* The DFS is extemely straightforward. I'll comment inline. */ void Solution::dfs(TreeNode *root, int r, int c) { if (root == nullptr) return; // If we're null, better return so you don't segfault. V[c][r].insert(root->val); // Insert the node's val into the proper place in the main data structure. dfs(root->left, r+1, c-1); // Recursively call the dfs on the left and right children. dfs(root->right, r+1, c+1); } /* The main method calls the dfs to set up the main data structure (V). Then it traverses that data structure to create the return vector. */ vector > Solution::verticalTraversal(TreeNode* root) { vector < vector > rv; map > >::iterator vit; map >::iterator mit; multiset ::iterator sit; int i; dfs(root, 0, 0); // Call the dfs on the root, to create V i = 0; // Traverse V to create the return value. for (vit = V.begin(); vit != V.end(); vit++) { rv.resize(i+1); for (mit = vit->second.begin(); mit != vit->second.end(); mit++) { for (sit = mit->second.begin(); sit != mit->second.end(); sit++) { rv[i].push_back(*sit); } } i++; } return rv; } /* Here's a main that reads nodes from top to bottom. If a value of -1 is specified, then the node is set to NULL. */ int main() { vector N; TreeNode *t; int n; int i, l, r, j; vector < vector > rv; Solution s; while (cin >> n) { if (n >= 0) { t = new TreeNode(n); N.push_back(t); } else { N.push_back(nullptr); } } for (i = 0; i < N.size(); i++) { if (N[i] != nullptr) { l = 2*i + 1; if (l < N.size()) N[i]->left = N[l]; l++; if (l < N.size()) N[i]->right = N[l]; } } rv = s.verticalTraversal(N[0]); for (i = 0; i < rv.size(); i++) { for (j = 0; j < rv[i].size(); j++) printf("%d ", rv[i][j]); printf("\n"); } return 0; }