在Python中查找树中距离恰好为k的不同顶点对的数量

假设我们有一个整数k，还有一个具有n个节点的树，我们必须计算距离恰好为k的不同顶点对的数量。

因此，如果输入如下：k = 2

则输出为4

为了解决这个问题，我们将遵循以下步骤：

  • N := 5005

  • graph := 大小为N的邻接表

  • vertex_count := 大小为505 x 5005的二维矩阵

  • res := 0

  • 定义一个函数insert_edge()。它将接收x, y

    • 将y插入graph[x]的末尾

    • 将x插入graph[y]的末尾

  • 定义一个函数dfs()。它将接收v, parent

  • vertex_count[v, 0] := 1

  • 对于graph[v]中的每个i，执行：

    • 如果i与parent不同，则

      • 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

更新于：2020年8月27日

195 次浏览

开启您的职业生涯

完成课程获得认证

开始学习
广告