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;
}

Download