好叶子节点对的数量(Java)

对于二叉树的任意两个不同的叶子节点 A、B，一个重要的性质是：它们之间有且仅有一条最短路径。设两个叶子节点的最近公共祖先为 P，则最短路径的长度，等于 A 到 P 的距离，加上 B 到 P 的距离。
于是，我们遍历所有非叶子节点 P，找到以 P 为最近公共祖先的所有叶子节点对，并根据上面的等式，计算每一对之间的距离，并统计距离不超过 distance 的节点对数目，就能够得到最终的答案。
那么怎么计算每个叶子节点对之间的距离呢？关键是要知道：以 P 为根节点的子树中的所有叶子节点，它们与 P 之间的距离。于是，对于任意的节点 P，我们先通过递归的方式，统计叶子节点与 P 的左孩子 left、右孩子 right 之间的距离；这样，两个以 P 为最近公共祖先的叶子节点 A、B，其中一个（例如 A）在以 left 为根的子树中，另一个（例如 B）在以 right 为根的子树中。A 与 B 之间的距离，就等于 A 与 left 之间的距离，加上 B 与 right 之间的距离，再加上 2。
由于本题中约束 distance≤10，因此对于每个非叶子节点 P，只需开辟长度为 distance+1 的数组 depths，其中 depths[i] 表示与 P 之间的距离为 i 的叶子节点数目。

给二叉树的根节点 root 和一个整数 distance 。
如果二叉树中两个 叶 节点之间的 最短路径长度 小于或者等于 distance ，那它们就可以构成一组 好叶子节点对 。
返回树中 好叶子节点对的数量 。

示例 1：

输入：root = [1,2,3,null,4], distance = 3
输出：1
解释：树的叶节点是 3 和 4 ，它们之间的最短路径的长度是 3 。这是唯一的好叶子节点对。

示例 2：

输入：root = [1,2,3,4,5,6,7], distance = 3
输出：2
解释：好叶子节点对为 [4,5] 和 [6,7] ，最短路径长度都是 2 。但是叶子节点对 [4,6] 不满足要求，因为它们之间的最短路径长度为 4 。

示例 3：
输入：root = [7,1,4,6,null,5,3,null,null,null,null,null,2], distance = 3
输出：1
解释：唯一的好叶子节点对是 [2,5] 。

示例 4：
输入：root = [100], distance = 1
输出：0

示例 5：
输入：root = [1,1,1], distance = 2
输出：1

提示：

tree 的节点数在 [1, 2^10] 范围内。
每个节点的值都在 [1, 100] 之间。
1 <= distance <= 10

package com.loo;

public class CountLeafPairs {

    public static void main(String[] args) {
        TreeNode root = new TreeNode(1);
        TreeNode left1 = new TreeNode(2);
        TreeNode right1 = new TreeNode(3);
        TreeNode left2 = new TreeNode(4);
        root.left = left1;
        root.right = right1;
        left1.right = left2;
        int distance = 3;
        System.out.println(getCountLeafPairs(root , distance));
    }
    
    public static int getCountLeafPairs(TreeNode root , int distance) {
        Pairs p = dfs(root , distance);
        return p.count;
    }
    
    public static Pairs dfs(TreeNode root , int distance) {
        int[] depths = new int[distance+1];
        boolean isLeaf = root.left == null && root.right == null;
        if (isLeaf) {
            depths[0] = 1;
            return new Pairs(depths , 0);
        }
        int[] leftDepths = new int[distance+1];
        int[] rightDepths = new int[distance+1];
        int leftCount = 0;
        int rightCount = 0;
        if (root.left!=null) {
            Pairs leftPairs = dfs(root.left , distance);
            leftDepths = leftPairs.depth;
            leftCount = leftPairs.count;
        }
        if (root.right!=null) {
            Pairs rightPairs = dfs(root.right , distance);
            rightDepths = rightPairs.depth;
            rightCount = rightPairs.count;
        }
        for (int i=0;i<distance;i++) {
            depths[i+1] += leftDepths[i];
            depths[i+1] += rightDepths[i];
        }
        int count = 0;
        for (int i=0;i<=distance;i++) {
            for (int j=0;j+i+2<=distance;j++) {
                count += leftDepths[i] * rightDepths[j];
            }
        }
        return new Pairs(depths , count + leftCount + rightCount);
    }
    
    static class TreeNode {
        int value;
        TreeNode left;
        TreeNode right;
        TreeNode(int v) {
            value = v;
        }
    }
    
    static class Pairs {
        int[] depth;
        int count;
        Pairs(int[] dep ,  int c) {
            depth = dep;
            count = c;
        }
    }

}