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:
https://gist.github.com/Rockbet/b4474ee277ad0ec5591725609dd6047f