#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); } }; } // namespace 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; } } // namespace _tests 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 */