MTREE - Another Tree Problem

Tác giả: happyboy99x

Ngôn ngữ: C++

#include<cstdio>
#include<vector>
using namespace std;

const int N = (int) 1e5 + 5, TWO_POW_MINUS_ONE = (int) 5e8 + 4;
const long long MOD = (int) 1e9 + 7;
#define TR(v,i) for(typeof((v).begin()) i=(v).begin();i!=(v).end();++i)

long long a[N], b[N];
int n, c[N]; bool vst[N];
vector<vector<pair<int, int> > > g;

void enter() {
	scanf("%d", &n); g.resize(n);
	for(int i = 1; i < n; ++i) {
		int u, v, l; scanf("%d%d%d", &u, &v, &l);
		g[--u].push_back(make_pair(--v, l));
		g[v].push_back(make_pair(u, l));
	}
}

void dfs(int u) {
	vst[u] = 1;
	TR(g[u], it) {
		int v = it->first, w = it->second;
		if(vst[v]) continue; dfs(v); long long tmp = w * (a[v] + 1) % MOD;
		a[u] = (a[u] + tmp) % MOD;
		b[u] = (b[u] - tmp * tmp % MOD) % MOD;
	}
	b[u] = (b[u] + a[u] * a[u] % MOD) % MOD * TWO_POW_MINUS_ONE % MOD;
}

void solve() {
	dfs(0); long long res = 0;
	for(int i = 0; i < n; ++i) res = (res + a[i] + b[i]) % MOD;
	res = (res + MOD) % MOD;
	printf("%d\n", (int) res);
}

int main() {
	enter();
	solve();
	return 0;
}

Download