golang 解法, 执行用时: 20 ms, 内存消耗: 9 MB, 提交时间: 2023-10-21 20:18:50

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func findDistance(root *TreeNode, p int, q int) int {
	var ans int
	var dfs func(root *TreeNode) (int, int)
	dfs = func(root *TreeNode) (int, int) {
		if root == nil {
			return -1, -1
		}
		pp, qq := -1, -1
		a, b := dfs(root.Left)
		if a != -1 {
			pp = a + 1
		}
		if b != -1 {
			qq = b + 1
		}
		a, b = dfs(root.Right)
		if a != -1 {
			pp = a + 1
		}
		if b != -1 {
			qq = b + 1
		}
		if root.Val == p {
			pp = 0
		}
		if root.Val == q {
			qq = 0
		}
		if pp != -1 && qq != -1 {
			ans = pp + qq
			pp,qq=-1,-1
		}
		return pp, qq
	}
	dfs(root)
	return ans
}

java 解法, 执行用时: 1 ms, 内存消耗: 43.6 MB, 提交时间: 2023-10-21 20:18:22

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public int findDistance(TreeNode root, int p, int q) {
        if(p == q){
            return 0;
        }
        // 公共祖先
        TreeNode parent = findCommonParent(root,p,q);
        // 节点到公共祖先的距离
        return getNodeDepth(parent,p) + getNodeDepth(parent,q);
    }

    // 首先寻找公共祖先
    private TreeNode findCommonParent(TreeNode root, int p, int q){
        if(root == null){
            return null;
        }
        if(root.val == p || root.val == q){
            return root;
        }
        TreeNode left = findCommonParent(root.left,p,q);
        TreeNode right = findCommonParent(root.right,p,q);
        if(left != null && right != null){
            return root;
        }
        return left == null ? right : left;
    }

    // 到root节点的距离
    private int getNodeDepth(TreeNode root,int p){
        if(root == null){
            // -1表示当前路径不是p的路径
            return -1;
        }
        if(root.val == p){
            return 0;
        }
        int left = getNodeDepth(root.left,p);
        int right = getNodeDepth(root.right,p);
        int ans = Math.max(left,right);
        // 最大值是-1的话，表示不是p的路径，>= 0则表示找到的p的路径，再加上自己本身的距离
        return ans >= 0 ? ans + 1 : -1;
    }

}

cpp 解法, 执行用时: 32 ms, 内存消耗: 30.9 MB, 提交时间: 2023-10-21 20:17:50

/**
 * Definition for a binary tree node.
 * 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 Solution {
public:
    int findDistance(TreeNode* root, int p, int q) 
    {
        TreeNode* LCA = find_LCA(root, p, q);       //先找LCA Least Common Ancestor
    
        int p_dep = dfs_calc_depth(LCA, p);
        int q_dep = dfs_calc_depth(LCA, q);
        return p_dep + q_dep;
    }

    TreeNode* find_LCA(TreeNode* root, int p, int q)
    {
        if (root==NULL || root->val == p || root->val == q)
            return root;
        TreeNode* L = find_LCA(root->left, p, q);
        TreeNode* R = find_LCA(root->right, p, q);

        if (L && R)
            return root;
        else if(L && R==NULL)
            return L;
        else if (L==NULL && R)
            return R;
        else
            return NULL;
    }

    int dfs_calc_depth(TreeNode* root, int target) 
    {
        if (root == NULL)
            return -1;
        if (root->val == target)
            return 0;
        int L = dfs_calc_depth(root->left, target);
        int R = dfs_calc_depth(root->right, target);
        if (L == -1 && R == -1)
            return -1;
        return max(L , R) + 1; 
    }
};

python3 解法, 执行用时: 76 ms, 内存消耗: 26.2 MB, 提交时间: 2023-10-21 20:17:19

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def findDistance(self, root: TreeNode, p: int, q: int) -> int:
        LCA = self.find_LCA(root, p, q)     #Least Common Ancestor
        
        p_dep = self.dfs(LCA, p)
        q_dep = self.dfs(LCA, q)
        return p_dep + q_dep

    def dfs(self, root: TreeNode, target: int) -> int:
        if root == None:
            return -1
        if root.val == target:
            return 0
        L = self.dfs(root.left, target)
        R = self.dfs(root.right, target)
        
        if L == -1 and R == -1:
            return -1
        return max(L , R) + 1

    def find_LCA(self, root: TreeNode, p: int, q: int) -> TreeNode:
        if root==None or root.val==p or root.val==q:
            return root
        L = self.find_LCA(root.left, p, q)
        R = self.find_LCA(root.right, p, q)
        if L and R:
            return root
        elif L and R==None:
            return L
        elif L == None and R:
            return R
        else:
            return None