Soluções Mapa

Primeiramente, note que o grafo é uma árvore. Se quisermos todos os caminhos que passam por determinado tipo de aresta, eliminamos todas essas arestas e nos restará uma floresta (várias componentes conexas e cada uma é uma árvore). Então, fica fácil ver que os caminhos que passam por um tipo de aresta são os caminhos entre vértices de componentes conexas distintas.

Seja $k$ o número de componentes conexas e $c_1$ , $c_2$ , ..., $c_k$ o número de vértices em cada uma delas. Queremos, então, a soma dos produtos dois a dois dos $c_i$ 's. Para fazer isso de maneira linear (quadrática estouraria o tempo!), usamos o fato de que:

$(c_1 + c_2 + ... + c_k)^2 = (c_1^2 + c_2^2 + ... + c_k^2) + 2(c_1 c_2 + c_1 c_3 + ... + c_2 c_3 + ... + c_{k-1} c_k)$

e podemos calcular os dois primeiros termos de maneira linear.

Para calcular os $c_i$ 's. podemos fazer um algoritmo de Flood Fill ou Union Find (vide as aulas do Curso Noic).

O código, então, fica:

	#include <cstdio>
	#include <vector>
	#include <cassert>
	#define MAXN 100100
	using namespace std;

	typedef long long ll;

	int n;
	bool visited[MAXN];
	vector<int> graph[MAXN];

	int last; //quantidade de componentes conexas
	ll qtt[MAXN];

	void dfs(int x){

	qtt[last]++;
	visited[x] = true;

	for(int i = 0;i < graph[x].size();i++){
	if(!visited[ graph[x][i] ])
	dfs(graph[x][i]);
	}
	}

	int main(){

	scanf("%d", &n);
	for(int i = 1;i < n;i++){
	int a, b, c;
	scanf("%d %d %d", &a, &b, &c);

	if(c) continue;

	graph[a].push_back(b);
	graph[b].push_back(a);
	}

	for(int i = 1;i <= n;i++) visited[i] = 0;
	for(int i = 1;i <= n;i++){
	if(!visited[i]){
	last++;
	dfs(i);
	}
	}

	ll answer = 0LL;
	for(int i = 1;i <= last;i++) answer += qtt[i];
	answer *= answer;
	for(int i = 1;i <= last;i++) answer -= qtt[i]*qtt[i];
	answer /= 2;

	printf("%lld\n", answer);

	return 0;
	}

view raw

mapa.cpp

hosted with ❤ by GitHub