MTREE - Another Tree Problem

Tác giả: hieult

Ngôn ngữ: C++

#include <vector>
#include <algorithm>
#include <utility>
#include <iostream>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <set>
#include <map>
#include <cstring>
#include <time.h>
#define rep(i,n) for(int i=0;i<n;i++)
#define pb push_back
#define Debug(x) cout<<#x<<"="<<x<<endl;
#define For(i,l,r) for(int i=l;i<=r;i++)
#define tr(e,x) for(eit e=x.begin();e!=x.end();e++)
const int inf=~0U>>1,maxn=100000,mod=1000000007;
using namespace std;
typedef long long ll;
int n;
struct Edge
{
    int t,c;
    Edge(int _t,int _c):t(_t),c(_c){}
};
vector<Edge> E[maxn];
typedef vector<Edge>::iterator eit;
void InsEdge(int s,int t,int c)
{
    E[s].pb(Edge(t,c));E[t].pb(Edge(s,c));
}
ll pow(int x,int e)
{
    if(!e)return 1;
    ll tmp=pow(x,e/2);tmp*=tmp;tmp%=mod;
    if(e&1)tmp*=x,tmp%=mod;
    return tmp;
}
ll Sum[maxn],P,ans=0;
int Q[maxn],F[maxn],h,t;
void BFS(int vs)
{
    h=t=0;
    for(Q[t++]=vs,F[vs]=-1;h<t;h++)
    {
        int x=Q[h];
        tr(e,E[x])if(e->t!=F[x])
            F[e->t]=x,Q[t++]=e->t;
    }
    for(int i=h-1;i>=0;i--)
    {
        int x=Q[i];Sum[x]=1;
        ll all,tmp,ret=0;
        tr(e,E[x])if(e->t!=F[x])
            Sum[x]+=Sum[e->t]*e->c,Sum[x]%=mod;
        all=Sum[x]+1;
        tr(e,E[x])if(e->t!=F[x])
            tmp=Sum[e->t]*e->c,tmp%=mod,ret+=(all-tmp)*tmp,ret%=mod;
        ret*=P;ret%=mod;if(ret<0)ret+=mod;ans+=ret;ans%=mod;
    }
}
int main()
{
    //freopen("in","r",stdin);
    cin>>n;int s,t,c;
    rep(i,n-1)
    {
        scanf("%d%d%d",&s,&t,&c);--s;--t;
        InsEdge(s,t,c);
    }
    P=pow(2,mod-2);
    BFS(0);
    cout<<ans<<endl;
}

Download