【Leetcode】1740. Find Distance in a Binary Tree

题目地址：

https://leetcode.com/problems/find-distance-in-a-binary-tree/

给定一棵二叉树，题目保证节点数值各不相同，再给定两个节点的值$p$和$q$，题目保证$p$和$q$在树中，问它们之间的距离。树的边权都视为$1$。

思路是分治。开一个DFS函数，其输入是树节点，并带上节点的深度作为参数，返回的是该子树中是否存在$p$与$q$其一，并且开三个全局变量，一个是最终答案，另外两个分别是$p$和$q$各自的深度。在DFS到当前节点的时候，如果当前节点是$p$或$q$，则给全局变量赋值一下深度。接着先递归求解左右子树，如果左右子树都含$p$与$q$其一，那么说明当前节点是$p$与$q$的最近公共祖先，那么它们的距离就是两者的深度减去当前节点深度的两倍；如果其一子树返回true，并且当前节点等于$p$或者$q$，那么距离就是两个深度的差，如果其一子树返回true，并且当前节点既不等于$p$也不等于$q$，则返回true给上层，因为找到了；否则说明两个子树都返回false，并且当前节点也不是$p$与$q$之一，则直接返回false给上层。代码如下：

public class Solution {
    
    private int res, pDepth, qDepth;
    
    public int findDistance(TreeNode root, int p, int q) {
        dfs(root, 0, p, q);
        return res;
    }
    
    // 返回root子树里是否存在p与q其一
    private boolean dfs(TreeNode root, int depth, int p, int q) {
        if (root == null) {
            return false;
        }
        
        if (root.val == p) {
            pDepth = depth;
        }
        if (root.val == q) {
            qDepth = depth;
        }
        
        // 在左右子树里递归寻找p和q
        boolean left = dfs(root.left, depth + 1, p, q), right = dfs(root.right, depth + 1, p, q);
        // 如果p与q分别在两个子树里，则距离就是深度和减去当前深度的两倍
        if (left && right) {
            res = pDepth + qDepth - 2 * depth;
            return true;
        }
        
        // 如果其一在当前节点，另一个在某个子树，则距离就是深度差
        if ((root.val == p || root.val == q) && (left || right)) {
            res = Math.abs(pDepth - qDepth);
            return true;
        }
        
        // 返回是否找到了其一
        return root.val == p || root.val == q || left || right;
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
    
    public TreeNode(int val) {
        this.val = val;
    }
}

时间复杂度$O(n)$，空间$O(h)$。