Vizualizare sursă pentru „4063 – Cartierul”


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

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

= Date de ieșire =
Programul va afișa pe ecran <code>n</code> numere, cel de al <code>i</code>-lea fiind numărul de culori distincte din subarborele cu rădăcina în <code>i</code>.

= Restricții și precizări =

* <code>1 ≤ n ≤ 200.000</code>
* Culorile arborelui sunt numere naturale mai mici sau egale cu <code>1.000.000.000</code>

= Exemplu: =
Intrare
 5
 2 3 2 2 1
 1 2
 1 3
 3 4
 3 5
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>