LeetCode/problems/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);
    }
};

}

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;
}

}

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 */