Блог пользователя RikiSabe

How to count the number of nodes with the same value for each node on a tree?
Автор RikiSabe, история, 8 месяцев назад, По-английски

Hi all, I have a question about queries on a tree.

You are given a tree and each node has a value, for each node 'u' you should count the number of nodes 'v' that exist in the subtree of node u with the same value than 'u'.

Input An integer n (n <= 10^5) the number of nodes in the tree Each of the next n — 1 lines describes and edge on the tree The last line consists of n integers the a_i (-10^5 <= a_i <= 10^5) the value for the node i (1, 2,... n)

Output n integers with the answer for the i-th node

EXAMPLE

INPUT:

9

1 3

3 5

3 4

1 2

2 6

2 7

7 8

7 9

5 2 -1 -1 5 5 -1 2 -1

OUTPUT

2 1 1 0 0 0 1 0 0

  • Проголосовать: нравится
  • 0
  • Проголосовать: не нравится

»
8 месяцев назад, # |
  Проголосовать: нравится +9 Проголосовать: не нравится

by using dsu on trees "small to large trick"

→ Ответить
  • »
    »
    8 месяцев назад, # ^ |
    ← Rev. 2   Проголосовать: нравится +3 Проголосовать: не нравится

    This is even simpler version and can be solved without any data structures in $$$O(n)$$$. 

    all(a) += 1e5;
freq = [0; 2e5];
dfs(u) {
  freq[a[u]] += 1;
  ans[u] -= freq[a[u]];
  for v in adj[u] {
    dfs(v);
  }
  ans[u] += freq[a[u]];
}
    → Ответить