Entendendo o Problema

Daddo uma árvore com n vértices e arestas coloridas (pretas ou vermelhas), precisamos contar sequências de k vértices que são consideradas "boas". Uma sequência é boa se, ao percorrer o caminho mais curto entre pares consecutivos de vértices, pelo menos uma aresta preta é utilizada durante todo o trajeto.

Análise do Problema

O segredo está em pensar inversamente: é mais fácil contar as sequências que não são boas e subtrair do total.

Uma sequência não é boa quando todos os caminhos percorridos utilizam apenas arestas vermelhas. Isso significa que todos os vértices da sequência pertencem ao mesmo componente conectado formado apenas por arestas vermelhas.

Se temos um componente de arestas vermelhas com s vértices, existem s^k sequências que permanecem inteiramente dentro desse componente — e essas sequências não são boas.

Portanto:

Sequências boas = n^k − Σ(s_i)^k

onde s_i representa o tamanho de cada componente de arestas vermelhas.

Solução Usando DFS

A primeira abordagem utiliza DFS para encontrar os componentes conectados formados apenas por arestas vermelhas.

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
const int LIMITE = 100000;
const int MOD = 1000000007;

int n, k;
vector<int> adj[LIMITE + 5];
bool procesado[LIMITE + 5];
ll respuesta, suma;

ll potencia(int base, int exponente) {
    ll resultado = 1;
    for (int i = 0; i < exponente; i++) {
        resultado = (resultado * base) % MOD;
    }
    return resultado;
}

void explorar(int vertice) {
    procesado[vertice] = true;
    for (int vizinhos : adj[vertice]) {
        if (!procesado[vizinhos]) {
            explorar(vizinhos);
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> n >> k;
    
    for (int i = 0; i < n - 1; i++) {
        int u, v, cor;
        cin >> u >> v >> cor;
        if (cor == 0) {
            adj[u].push_back(v);
            adj[v].push_back(u);
        }
    }
    
    respuesta = potencia(n, k);
    
    for (int i = 1; i <= n; i++) {
        if (!procesado[i]) {
            int tamanho = 0;
            function<void(int)> dfs = [&](int x) {
                procesado[x] = true;
                ++tamanho;
                for (int vizinho : adj[x]) {
                    if (!procesado[vizinho]) {
                        dfs(vizinho);
                    }
                }
            };
            dfs(i);
            suma = (suma + potencia(tamanho, k)) % MOD;
        }
    }
    
    cout << (respuesta - suma + MOD) % MOD << endl;
    return 0;
}

Solução Usando Union-Find

Uma alternativa mais elegante utiliza a estrutura de dados Union-Find (ou Disjoint Set Union) para agrupar os vértices conectados por arestas vremelhas.

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
const int LIMITE = 100000;
const int MOD = 1000000007;

int n, k;
int pai[LIMITE + 5];
int tamanhoComponente[LIMITE + 5];
bool marcado[LIMITE + 5];

ll potencia(int base) {
    ll resultado = 1;
    for (int i = 0; i < k; i++) {
        resultado = (resultado * base) % MOD;
    }
    return resultado;
}

int encontrar(int x) {
    if (pai[x] != x) {
        pai[x] = encontrar(pai[x]);
    }
    return pai[x];
}

void_unir(int x, int y) {
    int raizX = encontrar(x);
    int raizY = encontrar(y);
    if (raizX != raizY) {
        pai[raizX] = raizY;
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> n >> k;
    
    for (int i = 1; i <= n; i++) {
        pai[i] = i;
        tamanhoComponente[i] = 0;
    }
    
    for (int i = 0; i < n - 1; i++) {
        int u, v, cor;
        cin >> u >> v >> cor;
        if (cor == 0) {
            void_unir(u, v);
        }
    }
    
    ll total = potencia(n);
    
    for (int i = 1; i <= n; i++) {
        int raiz = encontrar(i);
        tamanhoComponente[raiz]++;
    }
    
    ll invalido = 0;
    for (int i = 1; i <= n; i++) {
        if (tamanhoComponente[i] > 0) {
            invalido = (invalido + potencia(tamanhoComponente[i])) % MOD;
        }
    }
    
    cout << (total - invalido + MOD) % MOD << endl;
    return 0;
}

Complexidade

• Tempo: O(n + k × número de componentes) ≈ O(n + k × log n) para Union-Find com path compression
• Memória: O(n) para armazenar a estrutura de adjacência ou os arrays do Union-Find

Exemplo de Execução

Para o primeiro exemplo de entrada:

4 4
1 2 1
2 3 1
3 4 1

Todas as arestas são pretas, então não há componentes de arestas vermelhas. Assim, todas as 4^4 = 256 sequências são boas, exceto as 4 sequências onde todos os elementos são iguais [1,1,1,1], [2,2,2,2], [3,3,3,3], [4,4,4,4]. O resultado é 252.