MTREE - Another Tree Problem
Tác giả: ladpro98
Ngôn ngữ: C++
#include <bits/stdc++.h>
#define ii pair<int, int>
#define X first
#define Y second
const int N = 100005;
const long long B = 1000000007;
using namespace std;
vector<ii> a[N], child[N];
bool was[N];
long long P[N], Q[N], res;
int n;
void DFS(int u) {
was[u] = true; int i, v;
for(i=0; i<a[u].size(); i++) {
v = a[u][i].X;
if (was[v]) continue;
child[u].push_back(a[u][i]);
DFS(v);
}
}
long long Pfunc(int u) {
if (P[u] != -1) return P[u];
P[u] = 0;
int i, v; long long uv;
for(i=0; i<child[u].size(); i++) {
v = child[u][i].X; uv = child[u][i].Y;
P[u] = (P[u] + (Pfunc(v) + 1) * uv) % B;
}
return P[u];
}
long long Qfunc(int u) {
if (Q[u] != -1) return Q[u];
int i, v;
long long aa = 0, bb = 0, uv, tmp = 0;
aa = Pfunc(u); aa = (aa * aa) % B;
for(i=0; i<child[u].size(); i++) {
v = child[u][i].X; uv = child[u][i].Y;
tmp = ((Pfunc(v) + 1) * uv) % B;
bb = (bb + tmp * tmp) % B;
}
tmp = (aa - bb + B) % B;
if (tmp % 2 == 1) tmp += B;
Q[u] = (tmp / 2) % B;
return Q[u];
}
int main()
{
scanf("%d\n", &n);
int i, u, v, w;
for(i=1; i<n; i++) {
scanf("%d %d %d\n", &u, &v, &w);
a[u].push_back(ii(v, w));
a[v].push_back(ii(u, w));
P[i] = Q[i] = -1;
}
P[n] = Q[n] = -1;
DFS(1);
long long res = 0;
for(i=1; i<=n; i++)
res = (res + Pfunc(i) + Qfunc(i)) % B;
printf("%lld", res);
return 0;
}