236. Lowest Common Ancestor of a Binary Tree

#include <iostream>
#include <vector>
#include <stack>
#include "TreeHelper.hpp"

using namespace std;

//        _______3______
//       /              \
//    ___5__          ___1__
//   /      \        /      \
//  6       2       0       8
// /  \
//7   4


//       37
//    /      \
// -34       -48
//     \    /    \
//  -100  -100    48
//                /
//              -54
//             /   \
//           -71    -22
//                    \
//                     8

class Solution {
private:
    bool p_found, q_found;

public:
    Solution() {
        p_found = false;
        q_found = false;
    }

    TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
        if (root == nullptr)
            return nullptr;
//        printf("node: %d\n", root->val);

        if (!p_found && root->val == p->val)
            p_found = true;
        if (!q_found && root->val == q->val)
            q_found = true;
        bool current_node_match = root->val == p->val || root->val == q->val;

        if (p_found && q_found && current_node_match)
            return root;

        //look in left sub-tree
        TreeNode *ret1 = lowestCommonAncestor(root->left, p, q);
        if (p_found && q_found)
            return current_node_match ? root : ret1;

        //look in right sub tree
        TreeNode *ret2 = lowestCommonAncestor(root->right, p, q);
        if (p_found && q_found)
            return current_node_match ? root : (ret1 && ret2 ? root : ret2);

        if (current_node_match)
            return root;
        return ret1 ? ret1 : ret2;
    }
};

class Tester {
public:
    static int test_num;

    template <class P1Type, class P2Type, class P3Type, class ExpType>
    void check(P1Type p1, P2Type p2, P3Type p3, ExpType expected) {
        Solution sol;
        test_num++;
        P1Type res = sol.lowestCommonAncestor(p1, p2, p3);
        if (res->val != expected) {
            cout << "test " << test_num << ": Failed: got " << res->val << endl;
        }
        else {
            printf("test %d: Passed\n", test_num);
        }
    }
};

int Tester::test_num = 0;


int main(int argc, const char * argv[]) {
    Solution sol;
    Tester t;

    string tree1 = "3,5,1,6,2,0,8,7,4,#,#,#,#,#,#,#,#,#,#";
    TreeNode *root = TreeHelper::deserialize<TreeNode>(tree1);

    t.check<TreeNode*, TreeNode*, TreeNode*, int>(root, root->left->left, root->left, 5);
    t.check<TreeNode*, TreeNode*, TreeNode*, int>(root, root->right->right, root->right, 1);
    t.check<TreeNode*, TreeNode*, TreeNode*, int>(root, root->left, root->right, 3);
    t.check<TreeNode*, TreeNode*, TreeNode*, int>(root, root->left->left->left, root->left->right, 5);

    string tree2 = "37,-34,-48,#,-100,-100,48,#,#,#,#,-54,#,-71,-22,#,#,#,8";
    root = TreeHelper::deserialize<TreeNode>(tree2);
    t.check<TreeNode*, TreeNode*, TreeNode*, int>(root, root->right->left, root->right->right->left->left, -48);

    return 0;
}