MTREE - Another Tree Problem

Tác giả: ll931110

Ngôn ngữ: C++

#include <cstring>
#include <iostream>
#include <vector>
#define mod 1000000007
using namespace std;

int pr[100010];
int n;
long long top[100010],bot[100010];
long long ret = 0;
vector< pair<int,int> > adj[100010];

void DFS(int u)
{
	for (vector< pair<int,int> >::iterator iter = adj[u].begin(); iter != adj[u].end(); iter++)
	{
		int v = iter->first;
		if (pr[v] > -1) continue;
		pr[v] = u;
		DFS(v);
	}
}

void DFScalc(int u)
{
	for (vector< pair<int,int> >::iterator iter = adj[u].begin(); iter != adj[u].end(); iter++)
	{
		int v = iter->first;
		if (pr[u] == v) continue;
		top[v] = (1LL * (top[u] + 1) * iter->second) % mod;
		DFScalc(v);
		long long tmp = (1LL * (bot[v] + 1) * iter->second) % mod;
		ret = (ret + bot[u] * tmp) % mod;
		bot[u] = (bot[u] + tmp) % mod;
	}
}

int main()
{
//	freopen("putevi.in","r",stdin);
//	freopen("putevi.out","w",stdout);
	scanf("%d", &n);
	for (int i = 1; i < n; i++)
	{
		int x,y,c;
		scanf("%d %d %d", &x, &y, &c);
		adj[x].push_back(make_pair(y,c));
		adj[y].push_back(make_pair(x,c));
	}
	memset(pr,-1,sizeof(pr));
	memset(top,0,sizeof(top));
	memset(bot,0,sizeof(bot));
	pr[1] = 0;
	DFS(1);
	DFScalc(1);
	for (int i = 1; i <= n; i++) ret = (ret + top[i]) % mod;
	cout << ret << endl;
}

Download