Solução Informática - Nível Avançado - Semana 3

Solução por Leonardo Paes

Conhecimento prévio necessário:

• Flood Fill (DFS)

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: