统计子树中城市之间最大距离

难度: Hard

给你 n 个城市，编号为从 1 到 n 。同时给你一个大小为 n-1 的数组 edges ，其中 edges[i] = [u_i, v_i] 表示城市 u_i 和 v_i之间有一条双向边。题目保证任意城市之间只有唯一的一条路径。换句话说，所有城市形成了一棵树。

一棵子树是城市的一个子集，且子集中任意城市之间可以通过子集中的其他城市和边到达。两个子树被认为不一样的条件是至少有一个城市在其中一棵子树中存在，但在另一棵子树中不存在。

对于 d 从 1 到 n-1 ，请你找到城市间 最大距离 恰好为 d 的所有子树数目。

请你返回一个大小为 n-1 的数组，其中第 d 个元素（下标从 1 开始）是城市间 最大距离 恰好等于 d 的子树数目。

请注意，两个城市间距离定义为它们之间需要经过的边的数目。

示例 1：

输入：n = 4, edges = [[1,2],[2,3],[2,4]]
输出：[3,4,0]
解释：
子树 {1,2}, {2,3} 和 {2,4} 最大距离都是 1 。
子树 {1,2,3}, {1,2,4}, {2,3,4} 和 {1,2,3,4} 最大距离都为 2 。
不存在城市间最大距离为 3 的子树。

示例 2：

输入：n = 2, edges = [[1,2]]
输出：[1]

示例 3：

输入：n = 3, edges = [[1,2],[2,3]]
输出：[2,1]

提示：

2 <= n <= 15
edges.length == n-1
edges[i].length == 2
1 <= u_i, v_i <= n
题目保证 (u_i, v_i) 所表示的边互不相同。

Submission

运行时间: 36 ms

内存: 16.1 MB

# class Solution:
#     def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
# 错了
#         s = set(range(1,n+1))
#         ans = [0]*(n-1)
#         def dfs(s):
#             res = 0
#             for x in s:
#                 for e in edges:
#                     if x in e:
#                         res = max(res,dfs(s-x)+1)
#                 res = max(res,)

# class Solution:
#     def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
#         g = [[] for _ in range(n)]
#         for [x,y] in edges:
#             g[x-1].append(y-1)
#             g[y-1].append(x-1)
#         in_set = [False] *n
#         ans = [0]*(n-1)
#         def f(i):
#             if i == n:
#                 diameter = 0
#                 vis = [False]*n
#                 for v,b in enumerate(in_set):
#                     if not b : continue
#                     def dfs(x):
#                         nonlocal diameter
#                         vis[x] = True
#                         max_len = 0
#                         for  y in g[x]:
#                             if not vis[y] and in_set[y]:
#                                 ml = dfs(y) +1
#                                 diameter = max(ml+max_len,diameter)
#                                 max_len = max(max_len,ml)
#                         return max_len
#                     dfs(v)
#                     break
#                 if diameter and vis == in_set:
#                     ans[diameter-1] +=1
#                 return

#             f(i+1)

#             in_set[i] = True
#             f(i+1)
#             in_set[i] = False
#             return
#         f(0)
#         return ans
class Solution:
    def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
        # 建树
        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x - 1].append(y - 1)
            g[y - 1].append(x - 1)  # 编号改为从 0 开始

        # 计算树上任意两点的距离
        dis = [[0] * n for _ in range(n)]
        def dfs(x: int, fa: int) -> None:
            for y in g[x]:
                if y != fa:
                    dis[i][y] = dis[i][x] + 1  # 自顶向下
                    dfs(y, x)
        for i in range(n):
            dfs(i, -1)  # 计算 i 到其余点的距离

        def dfs2(x: int, fa: int) -> int:
            # 能递归到这，说明 x 可以选
            cnt = 1  # 选 x
            for y in g[x]:
                if y != fa and \
                   (di[y] < d or di[y] == d and y > j) and \
                   (dj[y] < d or dj[y] == d and y > i):  # 满足这些条件就可以选
                    cnt *= dfs2(y, x)  # 每棵子树互相独立，采用乘法原理
            if di[x] + dj[x] > d:  # x 是可选点
                cnt += 1  # 不选 x
            return cnt
        ans = [0] * (n - 1)
        for i, di in enumerate(dis):
            for j in range(i + 1, n):
                dj = dis[j]
                d = di[j]
                ans[d - 1] += dfs2(i, -1)
        return ans

Explain

该题解采用状态压缩和二维 DP 的思路来解决问题。首先，通过 DFS 预处理计算出任意两点之间的距离，存储在二维数组 dis 中。然后，枚举任意两个点 i 和 j，通过 DFS 计算以 i 和 j 为直径端点的子树数量。在 DFS 过程中，对于每个点 x，如果满足一定条件（与 i 和 j 的距离关系），就可以选择 x，每棵子树之间互相独立，因此子树数量的计算可以采用乘法原理。最后，将不同直径对应的子树数量累加到答案数组 ans 中。

时间复杂度: O(n^3)

空间复杂度: O(n^2)

class Solution:
    def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
        # 建树
        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x - 1].append(y - 1)
            g[y - 1].append(x - 1)  # 编号改为从 0 开始

        # 计算树上任意两点的距离
        dis = [[0] * n for _ in range(n)]
        def dfs(x: int, fa: int) -> None:
            for y in g[x]:
                if y != fa:
                    dis[i][y] = dis[i][x] + 1  # 自顶向下
                    dfs(y, x)
        for i in range(n):
            dfs(i, -1)  # 计算 i 到其余点的距离

        def dfs2(x: int, fa: int) -> int:
            # 能递归到这，说明 x 可以选
            cnt = 1  # 选 x
            for y in g[x]:
                if y != fa and \
                   (di[y] < d or di[y] == d and y > j) and \
                   (dj[y] < d or dj[y] == d and y > i):  # 满足这些条件就可以选
                    cnt *= dfs2(y, x)  # 每棵子树互相独立，采用乘法原理
            if di[x] + dj[x] > d:  # x 是可选点
                cnt += 1  # 不选 x
            return cnt
        ans = [0] * (n - 1)
        for i, di in enumerate(dis):
            for j in range(i + 1, n):
                dj = dis[j]
                d = di[j]
                ans[d - 1] += dfs2(i, -1)
        return ans

Explore

在DFS预处理阶段，算法通过从每个节点i开始，使用DFS递归探索所有可达的节点，同时更新距离数组dis[i][y]。这种方式确保了从每个起点i到其他所有节点的距离都被探索并记录。由于树结构中任两节点间存在唯一路径，且DFS确保访问所有节点，因此不会存在未被正确更新的距离值。每次递归调用都传入父节点fa，防止重复访问和回溯，从而保证距离的准确性。

这些条件用于筛选出可以形成以i和j为直径端点的子树中的节点。条件`di[y] < d`或者`di[y] == d and y > j`确保节点y在i到j的路径上或者在i的方向上但不超过j。相似地，`dj[y] < d`或者`dj[y] == d and y > i`确保y在j的方向上但不超过i。这样的筛选条件是为了确保子树中的任何节点都直接或间接与i或j连接，符合子树的定义。然而，这种基于距离的筛选可能不会覆盖所有可能的情况，特别是在复杂的树结构中，可能需要更细致的逻辑来确保所有正确的子树被计算。

这一逻辑是基于直径的定义。如果`di[x] + dj[x] > d`，这说明节点x到i和j的距离之和大于i和j之间的距离，因此x位于以i和j为端点的路径之外。这种情况下，x可以作为一个独立的节点，不属于以i和j为直径的子树。因此，当出现这种情况时，需要额外增加计数来考虑x本身作为一个独立的子树节点的情形。这反映了在计算子树数量时，除了考虑明显的内部节点外，还需要考虑边界和外部节点的可能性。