tree_traversals.cc

#include <iostream>
#include <vector>
#include <stack>

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

class MyTreeOps {
    void trHelper(TreeNode* node, std::vector<int>& values, int type) {
        if (node) {
            switch (type) {
                case 0: // in-order
                    trHelper(node->left, values, type);
                    values.push_back(node->val);
                    trHelper(node->right, values, type);
                    break;
                case 1: // pre-order
                    values.push_back(node->val);
                    trHelper(node->left, values, type);
                    trHelper(node->right, values, type);
                    break;
                case 2: // post-order
                    trHelper(node->left, values, type);
                    trHelper(node->right, values, type);
                    values.push_back(node->val);
                default:
                    break;
            }
        }
    }
public:
    std::vector<int> traversalRecursion(TreeNode* root, int type) {
        std::vector<int> values;
        trHelper(root, values, type);
        return values;
    }

    std::vector<int> inorderTraversal(TreeNode* root) {
        TreeNode* curr = root;
        std::vector<int> values;
        std::stack<TreeNode*> s;

        while (curr || !s.empty()) {
            if (curr) {
                s.push(curr);
                curr = curr->left;
            } else {
                curr = s.top();
                values.push_back(curr->val);
                curr = curr->right;
                s.pop();
            }
        }

        return values;
    }

    std::vector<int> preorderTraversal(TreeNode* root) {
        TreeNode* curr = root;
        std::vector<int> values;
        std::stack<TreeNode*> s;

        s.push(curr);

        while(!s.empty()) {
            curr = s.top();
            s.pop();
            values.push_back(curr->val);

            if (curr->right)
                s.push(curr->right);

            if (curr->left)
                s.push(curr->left);
        }

        return values;
    }

    std::vector<int> postorderTraversal(TreeNode* root) {
        TreeNode* curr = root;
        std::vector<int> values;
        std::stack<TreeNode*> s;

        s.push(curr);

        while(!s.empty()) {
            if (curr->right)
                s.push(curr->right);

            if (curr->left)
                s.push(curr->left);

            curr = s.top();
            s.pop();
            values.push_back(curr->val);
        }

        return values;
    }

    // Morris traversal (Threading), O(n) time and O(1) space, in-order only
    // https://en.wikipedia.org/wiki/Threaded_binary_tree
    std::vector<int> inorderTraversalMorris(TreeNode* root) {
        TreeNode* curr = root;
        TreeNode* pred;
        std::vector<int> nodes;

        while (curr) {
            if (curr -> left) {
                pred = curr -> left;

                while (pred -> right && pred -> right != curr)
                    pred = pred -> right;

                if (!pred -> right) {
                    pred -> right = curr;
                    curr = curr -> left;
                } else  {
                    pred -> right = nullptr;
                    nodes.push_back(curr -> val);
                    curr = curr -> right;
                }
            } else {
                nodes.push_back(curr -> val);
                curr = curr -> right;
            }
        }

        return nodes;
    }
};

int main() {
    TreeNode a = TreeNode(1);
    TreeNode b = TreeNode(3);
    TreeNode c = TreeNode(2, &a, &b);

    MyTreeOps t = MyTreeOps();

    for (int val : t.traversalRecursion(&c, 0))
        std::cout << val << "\n";

    std::cout << "\n";

    for (int val : t.inorderTraversalMorris(&c))
        std::cout << val << "\n";

    return 0;
}