Python의 트리에서 거리가 정확히 k인 고유한 정점 쌍의 수 찾기

정수 k가 있고 n개의 노드가 있는 트리도 있다고 가정하면 정확한 k 거리를 갖는 고유한 정점 쌍의 수를 계산해야 합니다.

따라서 입력이 k =2

와 같으면

그러면 출력은 4가 됩니다.

이 문제를 해결하기 위해 다음 단계를 따릅니다. −

N :=5005
graph :=N 사이즈의 인접 리스트
vertex_count :=505 x 5005 크기의 2차원 행렬
해상도 :=0
insert_edge() 함수를 정의합니다. x, y
가 걸립니다.
- 그래프[x]의 끝에 y를 삽입
- 그래프 끝에 x 삽입[y]
dfs() 함수를 정의합니다. 이것은 v, 부모가 필요합니다.
vertex_count[v, 0] :=1
그래프[v]의 각 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)