Chương trình đếm số node trong subtree có cùng nhãn bằng Python

Cho cây gốc n node (0..n-1), mỗi node có nhãn chữ thường. labels[i] là nhãn node i. Cây cho bằng edge list [u,v] (u cha, v con). Tìm mảng A kích thước n: A[i] = số node trong subtree của i có nhãn giống i.

Ví dụ: n=5, edges=[[0,1],[0,2],[1,3],[0,4]], labels="ccaca" → [3,2,1,1,1].

Thuật toán

• Xây graph từ edge list (defaultdict set)
• N = {i: label} mapping node→label
• R = [0]*n kết quả
• DFS r(ni):
  • C = Counter()
  • Với con e của ni: xóa ngược边, C.update(r(e))
  • C.update(N[ni])
  • R[ni] = C[N[ni]]
  • Return C
• Gọi r(0), return R

Code Python

from collections import defaultdict, Counter

def solve(n, edges, labels):
    E = defaultdict(set)
    for f, t in edges:
        E[f].add(t)
        E[t].add(f)
    N = {i: e for i, e in enumerate(labels)}
    R = [0] * n
    def r(ni):
        C = Counter()
        for e in E[ni]:
            E[e].remove(ni)
            C.update(r(e))
        C.update(N[ni])
        R[ni] = C[N[ni]]
        return C
    r(0)
    return R

n = 5
edges = [[0,1],[0,2],[1,3],[0,4]]
labels = "ccaca"
print(solve(n, edges, labels))

Input

5, [[0,1],[0,2],[1,3],[0,4]], "ccaca"

Output

[3, 2, 1, 1, 1]