Solução por Leonardo Paes
Conhecimento prévio necessário:
Para resolvermos esse problema, utilizaremos uma DFS para percorrer a árvore e precalcular o tamanho de cada subárvore, considerando que a árvore está enraizada no vértice de número 1.
Tendo o tamanho de cada subárvore calculado, faremos o seguinte: Suponha que estamos no vértice $$u$$, com cor $$C_u$$. Iremos considerar nesse passo apenas os caminhos que passam através de $$u$$ sem passar por vértices de cor $$C_u$$ previamente calculados na ordem da $$DFS$$, ou seja, iremos aumentar apenas caminhos de cor $$C_u$$. Temos dois tipos de caminhos a considerar: caminhos que começam em alguma subárvore de $$u$$, passam por $$u$$ e terminam em outra subárvore de $$u$$. E caminhos que começam em alguma subárvore de $$u$$ e terminam em algum vértice que não está na subárvore de $$u$$.
Seja $$k$$ a quantidade de filhos de $$u$$ e seja $$sz_i$$ o tamanho da subárvore do i-ésimo filho. A quantidade de caminhos do primeiro tipo é igual ao seguinte somatório:
$$\sum_{i=1}^{k-1} \Big ( \sum_{j=i+1}^k sz_i \cdot sz_j \Big )$$
Para calcularmos isso em um tempo menor que $$O(n^2)$$, basta observarmos que: $$sz_1$$ será multiplicado com $$(sz_2 + sz_3 + … + sz_k)$$, $$sz_2$$ será multiplicado com $$(sz_3 + sz_4 + … + sz_k)$$, e assim por diante. Portanto, sendo $$sum = sz_1 + sz_2 + sz_3 + … + sz_k$$, basta utilizarmos o seguinte algoritmo, em pseudo-código:
Para cada filho $$i$$ de $$u$$, retiramos seu tamanho, $$sz_i$$ de $$sum$$ e incrementamos a resposta de $$C_u$$ em $$sz_i * sum$$.
Para calcularmos os caminhos do segundo tipo, devemos tomar muito cuidado, para evitar recontar caminhos diversas vezes, pois devemos contá-los apenas uma vez.
Seja $$used[C_u]$$ a quantidade de vértices tal que, partindo de um vértice $$u$$ de cor $$C_u$$, eu chego nesses vértices após passar por algum outro vértice de cor $$C_u$$, sem ser $$u$$, que foram calculados antes na ordem da DFS. Devemos evitar contar caminhos que terminam nesses vértices, pois eles já foram contados previamente. Portanto, temos dois casos a considerar para calcular o valor de $$used[C_u]$$: ao chamar algum filho de $$u$$ para ser calculado e ao terminar de calcular a resposta para toda a subárvore de $$u$$.
Antes de tudo, criaremos uma variável $$previous$$_$$used$$ que guardará o $$used[C_u]$$ inicial.
No primeiro caso, ao chamarmos algum filho $$v$$ para ser calculado, devemos considerar que toda a árvore, exceto a subárvore desse filho, já foi calculada. Portanto, $$used[C_u] = n – sz_v$$.
No segundo caso, após terminarmos de calcular a subárvore de $$u$$, temos que $$used[C_u] = previous$$_$$used + sz_u$$, pois, para todos os vértices que posteriormente serão calculados que tem a cor $$C_u$$, para chegar em algum filho $$v$$ de $$u$$, será necessário passar por $$u$$.
Após termos calculado o $$used$$ corretamente, ao calcularmos a resposta de um vértice $$u$$, incrementamos a resposta da cor $$C_u$$ em $$(n – sz[u] – used[C_u] + 1) * sz[u]$$, pois contaremos os caminhos começando em alguém da subárvore de $$u$$ e terminando em alguém fora dela, sem passar novamente em algum vértice de cor $$C_u$$.
Confira o código abaixo para melhor entendimento:
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #include <bits/stdc++.h> | |
| using namespace std; | |
| const int maxn = 2e5+10; | |
| typedef long long ll; | |
| int n, c[maxn], sz[maxn], used[maxn]; | |
| ll ans[maxn]; | |
| vector<int> grafo[maxn]; | |
| void dfs(int u, int p = 0){ | |
| sz[u] = 1; | |
| for(auto v : grafo[u]){ | |
| if(v == p) continue; | |
| dfs(v, u); | |
| sz[u] += sz[v]; | |
| } | |
| } | |
| void solve(int u, int p = 0){ | |
| ans[c[u]] += 1LL*(n – sz[u] – used[c[u]] + 1) * sz[u]; | |
| ll sum = sz[u] – 1; | |
| for(auto v : grafo[u]){ | |
| if(v == p) continue; | |
| sum -= sz[v]; | |
| ans[c[u]] += sz[v] * sum; | |
| } | |
| int previous_used = used[c[u]]; | |
| for(auto v : grafo[u]){ | |
| if(v == p) continue; | |
| used[c[u]] = n – sz[v]; | |
| solve(v, u); | |
| } | |
| used[c[u]] = previous_used + sz[u]; | |
| } | |
| int main(){ | |
| ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL); | |
| cin >> n; | |
| for(int i=1; i<=n; i++) cin >> c[i]; | |
| for(int i=1; i<n; i++){ | |
| int u, v; | |
| cin >> u >> v; | |
| grafo[u].push_back(v); | |
| grafo[v].push_back(u); | |
| } | |
| dfs(1); | |
| solve(1); | |
| for(int i=1; i<=n; i++) cout << ans[i] << "\n"; | |
| return 0; | |
| } |