在 Python 中查找树中距离恰好为 k 的不同顶点对的数量
pythonserver side programmingprogramming更新于 2025/10/15 19:24:17
假设我们有一个整数 k，并且有一棵有 n 个节点的树，我们必须计算距离恰好为 k 的不同顶点对的数量。
因此，如果输入为 k = 2
则输出将为 4
要解决这个问题，我们将遵循以下步骤 −
N := 5005
graph :=大小为 N 的邻接列表
vertex_count := 大小为 505 x 5005 的 2d 矩阵
res := 0
定义一个函数 insert_edge() 。这将需要 x、y
在 graph[x] 的末尾插入 y
在 graph[y] 的末尾插入 x
定义一个函数 dfs() 。这将采用 v，父节点
vertex_count[v, 0] := 1
对于 graph[v] 中的每个 i，执行
如果 i 与父节点不同，则
dfs(i, v)
对于范围从 1 到 k + 1 的 j，执行
res := res + vertex_count[i, j - 1] * vertex_count[v, k - j]
对于范围从 1 到 k + 1 的 j，执行
vertex_count[v, j] := vertex_count[v, j] + vertex_count[i, j - 1]
示例
让我们看看下面的实现以便更好地理解 −
N = 5005 graph = [[] for i in range(N)] vertex_count = [[0 for i in range(505)] for i in range(N)] res = 0 def insert_edge(x, y): graph[x].append(y) graph[y].append(x) def dfs(v, parent): global res vertex_count[v][0] = 1 for i in graph[v]: if (i != parent): dfs(i, v) for j in range(1, k + 1): res += vertex_count[i][j - 1] * vertex_count[v][k - j] for j in range(1, k + 1): vertex_count[v][j] += vertex_count[i][j - 1] k = 2 insert_edge(1, 2) insert_edge(2, 3) insert_edge(3, 4) insert_edge(2, 5) dfs(1, 0) print(res)
输入
k = 2 insert_edge(1, 2) insert_edge(2, 3) insert_edge(3, 4) insert_edge(2, 5)
输出
4
相关文章
有用资源
python 参考教程 - 该教程包含有关 python 的更多信息:https://www.cainiaomax.com/python/
