Editing 4229 – Kdist (section)

== Rezolvare ==
<syntaxhighlight lang="python3">
from collections import defaultdict, deque


def calculate_color_distances(n, k, edges, colors):
 
    tree = defaultdict(list)
    for u, v in edges:
        tree[u].append(v)
        tree[v].append(u)


    color_nodes = defaultdict(list)
    for i in range(n):
        color_nodes[colors[i]].append(i + 1)


    def bfs_distance_sum(start, valid_nodes):
        visited = set()
        queue = deque([(start, 0)])
        visited.add(start)
        distance_sum = 0
        node_count = 0

        while queue:
            node, dist = queue.popleft()
            distance_sum += dist
            node_count += 1

            for neighbor in tree[node]:
                if neighbor not in visited and neighbor in valid_nodes:
                    visited.add(neighbor)
                    queue.append((neighbor, dist + 1))

        return distance_sum, node_count


    color_distance_sums = [0] * k
    for color in range(1, k + 1):
        nodes = color_nodes[color]
        if len(nodes) < 2:
            continue

        total_distance_sum = 0
        for node in nodes:
            dist_sum, node_count = bfs_distance_sum(node, set(nodes))
            total_distance_sum += dist_sum


        color_distance_sums[color - 1] = total_distance_sum // 2

    return color_distance_sums



n, k = map(int, input().split())
edges = [tuple(map(int, input().split())) for _ in range(n - 1)]
colors = [int(input()) for _ in range(n)]


result = calculate_color_distances(n, k, edges, colors)
for dist_sum in result:
    print(dist_sum)

</syntaxhighlight>