#include #include #include // Definition for a binary tree node. 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) { } }; namespace { struct TreeNodeHandle { TreeNode* node; int x; int y; TreeNodeHandle(TreeNode* node, int x, int y) : node(node) , x(x) , y(y) { } bool operator==(const TreeNodeHandle& rhs) const { return node == rhs.node && x == rhs.x && y == rhs.y; } int value() const { return node->val; } bool operator<(const TreeNodeHandle& rhs) const { if (y != rhs.y) { return y < rhs.y; } if (x != rhs.x) { return x < rhs.x; } return value() < rhs.value(); } }; class verticals { TreeNode* root; class verticals_iter { std::deque queue; void advance() { auto& n = queue.front(); if (n.node->left != nullptr) { queue.push_back(TreeNodeHandle(n.node->left, n.x - 1, n.y + 1)); } if (n.node->right != nullptr) { queue.push_back(TreeNodeHandle(n.node->right, n.x + 1, n.y + 1)); } } public: verticals_iter(std::deque queue) : queue(queue) { if (queue.front().node == nullptr) { queue.pop_front(); } } bool operator!=(const verticals_iter& other) const { return queue != other.queue; } verticals_iter operator++() { advance(); queue.pop_front(); return *this; } TreeNodeHandle& operator*() { return queue.front(); } }; public: verticals(TreeNode* root) : root(root) { } verticals_iter begin() const { return verticals_iter(std::deque { TreeNodeHandle(root, 0, 0) }); } verticals_iter end() const { std::deque q; return verticals_iter(q); } }; } class Solution { public: std::vector> verticalTraversal(TreeNode* root) { std::map> traversals; int min_x = 0; int max_x = 0; for (auto node : verticals(root)) { traversals[node.x].push_back(node); min_x = std::min(min_x, node.x); max_x = std::max(max_x, node.x); } std::vector> result; for (int x = min_x; x <= max_x; x++) { auto& v = traversals[x]; if (v.size()) { std::sort(v.begin(), v.end()); result.push_back(std::vector {}); for (auto& n : v) { result.back().push_back(n.value()); } } } return result; } }; #pragma region tests #include namespace _tests { TreeNode* construct_tree(const std::vector& values, std::size_t i = 0) { if (i >= values.size() || values[i] == -1) { return nullptr; } auto tree = new TreeNode(values[i]); tree->left = construct_tree(values, 2 * i + 1); tree->right = construct_tree(values, 2 * i + 2); return tree; } void destruct_tree(TreeNode* ptr) { if (ptr == nullptr) { return; } destruct_tree(ptr->left); destruct_tree(ptr->right); delete ptr; } } TEST(examples, first) { Solution s; auto t = _tests::construct_tree(std::vector { 3, 9, 20, -1, -1, 15, 7 }); EXPECT_EQ(s.verticalTraversal(t), (std::vector { std::vector { 9 }, std::vector { 3, 15 }, std::vector { 20 }, std::vector { 7 } })); _tests::destruct_tree(t); } TEST(examples, second) { Solution s; auto t = _tests::construct_tree(std::vector { 1, 2, 3, 4, 5, 6, 7 }); EXPECT_EQ(s.verticalTraversal(t), (std::vector { std::vector { 4 }, std::vector { 2 }, std::vector { 1, 5, 6 }, std::vector { 3 }, std::vector { 7 } })); _tests::destruct_tree(t); } TEST(examples, third) { Solution s; auto t = _tests::construct_tree(std::vector { 1, 2, 3, 4, 6, 5, 7 }); EXPECT_EQ(s.verticalTraversal(t), (std::vector { std::vector { 4 }, std::vector { 2 }, std::vector { 1, 5, 6 }, std::vector { 3 }, std::vector { 7 } })); _tests::destruct_tree(t); } TEST(submission, first) { Solution s; auto t = _tests::construct_tree(std::vector { 3, 1, 4, 0, 2, 2 }); EXPECT_EQ(s.verticalTraversal(t), (std::vector { std::vector { 0 }, std::vector { 1 }, std::vector { 3, 2, 2 }, std::vector { 4 } })); _tests::destruct_tree(t); } int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); } #pragma endregion /* tests */