Số cặp đỉnh khác nhau có khoảng cách chính xác k trong cây trong Python

Cho số nguyên k và một cây n nút. Đếm số cặp đỉnh khác nhau có khoảng cách chính xác k.

Ví dụ: k=2 với cây như hình -> kết quả 4.

Thuật toán giải quyết (DP trên cây)

• N := 5005, graph danh sách kề.
• vertex_count[v][d]: số nút trong cây con gốc v ở độ sâu d.
• DFS từ gốc:
  • vertex_count[v][0] = 1.
  • Với mỗi con i:
    • Gọi DFS(i).
    • Với j từ 1 đến k: res += vertex_count[i][j-1] * vertex_count[v][k-j].
    • Cập nhật vertex_count[v][j] += vertex_count[i][j-1].
• Trả về res.

Ví dụ minh họa

N = 5005
graph = [[] for _ in range(N)]
vertex_count = [[0]*505 for _ 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)  # 4

Đầu vào

k=2, edges: (1,2),(2,3),(3,4),(2,5)

Đầu ra

4