Cousins in Binary Tree

• Time : $O(n)$
• Space : $O(2^h_B)$

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */

void dfs1(TreeNode* A, int t, int p, int dep, int& level, int& par) {
    if(!A) return;
    if(A->val == t) {
        par = p;
        level = dep;
    }
    dfs1(A->left,t,A->val,dep + 1, level, par);
    dfs1(A->right,t,A->val,dep + 1, level, par);
}

void dfs2(TreeNode* A, int p, int dep, int& level, int& par, vector<int>& res) {
    if(!A) return;
    if(dep > level) return;
    if(dep == level) {
        if(p != par) res.push_back(A->val);
        return;
    }
    dfs2(A->left,A->val,dep + 1, level, par, res);
    dfs2(A->right,A->val,dep + 1, level, par, res);
}

vector<int> Solution::solve(TreeNode* A, int B) {
    if(!A) return {};
    if(A->val == B) return {};
    int level = -1, par = -1;
    dfs1(A,B,-1,0,level,par);
    vector<int> res;
    dfs2(A,-1,0,level,par,res);
    return res;
}