Cerința

Avem în cartier un arbore cu nodurile colorate și cu rădăcina nodul 1. Pentru fiecare nod i al arborelui să se afișeze câte culori distincte sunt în subarborele cu rădăcina în i.

Date de intrare

Programul citește de la tastatură numărul n, reprezentând numărul de noduri al arborelui. Pe a doua linie se vor afla n valori, cea de a i reprezentând culoarea nodului i. Urmează n-1 linii, pe fiecare linie aflânduse 2 numere a b indicând faptul ca între a și b este muchie.

Date de ieșire

Programul va afișa pe ecran n numere, cel de al i-lea fiind numărul de culori distincte din subarborele cu rădăcina în i.

Restricții și precizări

1 ≤ n ≤ 200.000
Culorile arborelui sunt numere naturale mai mici sau egale cu 1.000.000.000

Exemplu:

Intrare

Ieșire

3 1 2 1 1

Rezovlare

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

def distinct_colors_in_subtrees(n, colors, edges):

   tree = defaultdict(list)
   for a, b in edges:
       tree[a].append(b)
       tree[b].append(a)

   subtree_colors = [set() for _ in range(n + 1)]
   result = [0] * (n + 1)
   visited = [False] * (n + 1)

   def dfs(node):
       visited[node] = True
       subtree_colors[node].add(colors[node - 1])  # Add the color of the current node
       for neighbor in tree[node]:
           if not visited[neighbor]:
               dfs(neighbor)
               # Merge the color sets of the subtrees
               if len(subtree_colors[neighbor]) > len(subtree_colors[node]):
                   subtree_colors[node], subtree_colors[neighbor] = subtree_colors[neighbor], subtree_colors[node]
               subtree_colors[node].update(subtree_colors[neighbor])
       result[node] = len(subtree_colors[node])

   dfs(1)

   return result[1:]

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

result = distinct_colors_in_subtrees(n, colors, edges)

print("\n".join(map(str, result)))

</syntaxhighlight>