LeetCode/cpp/vertical-order-traversal-of-a-binary-tree.cpp

#include <deque>
#include <map>
#include <vector>

// 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<TreeNodeHandle> 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<TreeNodeHandle> 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<TreeNodeHandle> q;
        return verticals_iter(q);
    }
};

} // namespace

class Solution {
  public:
    std::vector<std::vector<int>> verticalTraversal(TreeNode *root) {
        std::map<int, std::vector<TreeNodeHandle>> 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<std::vector<int>> 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<int>{});

                for (auto &n : v) {
                    result.back().push_back(n.value());
                }
            }
        }

        return result;
    }
};

#pragma region tests

#include <gtest/gtest.h>

namespace _tests {

TreeNode *construct_tree(const std::vector<int> &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 */