Binary Tree - Divide Conquer & Traverse

public class Solution {
    /**
     * @param root: the given BST
     * @param k: the given k
     * @return: the kth smallest element in BST
     */
    public int kthSmallest(TreeNode root, int k) {
        // write your code here
        Map<TreeNode, Integer> numOfChildren = new HashMap<>();
        countNodes(root, numOfChildren);
        return quickSelectOnTree(numOfChildren, k, root);
    }

    public int countNodes(TreeNode root, Map<TreeNode, Integer> numOfChildren) {
        if(root == null) {
            return 0;
        }
        int left = countNodes(root.left, numOfChildren);
        int right = countNodes(root.right, numOfChildren);
        numOfChildren.put(root, left + right + 1);
        return left + right + 1;
    }
    
    public int quickSelectOnTree(Map<TreeNode, Integer> numOfChildren, int k, TreeNode root) {
        if(root == null) {
            return -1;
        }
        int left = root.left == null ? 0 : numOfChildren.get(root.left);
        if(left >= k) {
            return quickSelectOnTree(numOfChildren, k, root.left);
        }
        if(left + 1 == k) {
            return root.val;
        }
        return quickSelectOnTree(numOfChildren, k - left - 1, root.right);
    }
}

/**
 * Definition of TreeNode:
 * public class TreeNode {
 *     public int val;
 *     public TreeNode left, right;
 *     public TreeNode(int val) {
 *         this.val = val;
 *         this.left = this.right = null;
 *     }
 * }
 */
class ResultType {
    public boolean a_exist, b_exist;
    public TreeNode node;
    ResultType(boolean a, boolean b, TreeNode n) {
        a_exist = a;
        b_exist = b;
        node = n;
    }
}

public class Solution {
    /*
     * @param root: The root of the binary tree.
     * @param A: A TreeNode
     * @param B: A TreeNode
     * @return: Return the LCA of the two nodes.
     */
    public TreeNode lowestCommonAncestor3(TreeNode root, TreeNode A, TreeNode B) {
        // write your code here
        ResultType rt = helper(root, A, B);
        if(rt.a_exist && rt.b_exist) {
            return rt.node;
        } else {
            return null;
        }
    }
    
    public ResultType helper(TreeNode root, TreeNode A, TreeNode B) {
        if(root == null) {
            return new ResultType(false, false, null);
        }
        
        ResultType left_rt = helper(root.left, A, B);
        ResultType right_rt = helper(root.right, A, B);
        
        boolean a_exist = left_rt.a_exist || right_rt.a_exist || root == A;
        boolean b_exist = left_rt.b_exist || right_rt.b_exist || root == B;
        
        if(root == A || root == B) {
            return new ResultType(a_exist, b_exist, root);
        }
        
        if(left_rt.node != null && right_rt.node != null) {
            return new ResultType(a_exist, b_exist, root);
        }
        if(left_rt.node != null) {
            return new ResultType(a_exist, b_exist, left_rt.node);
        }
        if(right_rt.node != null) {
            return new ResultType(a_exist, b_exist, right_rt.node);
        }
        return new ResultType(a_exist, b_exist, null);
    }
}