题目描述

给定一棵 N 叉树，从树根开始，找出所有与其所有子树中的叶节点距离相等的节点数量。

示例

例如：

           0
   /   /     \   \
  1   2       3   4           0
     / | \    |   \
    5  6  7   8    9          |
                             -----
                             | 2 |
                             -----

在这个示例中，距离根节点 0 为 0 的节点有根节点 0自己。距离为 1 的节点有 1,2,3,4。距离为 2 的节点有一个节点 5。距离为 3 的节点有一个节点 6。距离为 4 的节点有三个节点 7,8,9。不存在距离为 5 的节点。和为所有节点数量的 10，但是只有 2 个节点符合要求（距离为 2 的节点 5 和距离为 3 的节点 6）。

解法

本题可以使用树遍历中的后序遍历，从下往上统计节点的子树内所有叶节点到该节点的距离，同时记录每个距离对应的节点数量，最终返回所有距离为相同值的节点数量之和，即为所求的结果。

同时，在遍历过程中也需要记录每个节点到所有子节点的距离，以便进行从下往上的距离计算。

具体实现可以通过递归函数完成。

代码实现

Python实现代码示例：

class Solution:
    def countSubTree(self, n: int, edges: List[List[int]]) -> List[int]:
        # 建立邻接表
        adj = [[] for _ in range(n)]
        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)

        # 后序遍历
        def dfs(node, parent):
            count = 0  # 该节点子树内所有叶节点到该节点的距离之和
            num = {}   # 记录该距离下的节点数量
            for child in adj[node]:
                if child == parent:
                    continue
                child_count, child_num = dfs(child, node)
                for d, c in child_num.items():
                    if d + 1 not in num:
                        num[d + 1] = 0
                    num[d + 1] += c
                count += child_count + sum(child_num.values())
            if not num:
                num[0] = 1
            for d, c in num.items():
                res[d] += c * (count - sum(num[i] for i in range(d)))
            return count, num

        res = [0] * n
        dfs(0, -1)
        return res

其中，输入参数 n 表示节点总数，edges 表示节点对之间的边，这里使用邻接表存储树结构。返回值 res 表示每个节点的子树内所有叶节点到该节点的距离相等的节点数量。