Editing 4229 – Kdist (section)

= Exemplu: =
Intrare
 6 3
 1 2
 1 3
 3 4
 3 5
 5 6
 1
 2
 2
 1
 2
 3
Ieșire
 2
 6
 0

=== Explicație ===
Pentru culoarea <code>1</code> avem o pereche: <code>(1, 4)</code>, cu distanţa <code>2</code>. Pentru culoarea <code>2</code> avem trei perechi: <code>(2, 3)</code> cu distanţa <code>2</code>, <code>(2, 5)</code> cu distanţa <code>3</code> și <code>(3, 5)</code> cu distanţa <code>1</code>. Pentru culoarea <code>3</code> avem un singur nod cu această culoare, deci răspunsul este <code>0</code>.\

== 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>