1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
| #include <algorithm> #include <iostream> #include <cstring> #include <cstdio> #include <vector> #include <stack> #include <queue> #include <cmath> #include <set> #include <map> #define mp make_pair #define pb push_back #define pii pair<int,int> #define link(x) for(edge *j=h[x];j;j=j->next) #define inc(i,l,r) for(int i=l;i<=r;i++) #define dec(i,r,l) for(int i=r;i>=l;i--) const int MAXN=3e5+10; const double eps=1e-8; #define ll long long using namespace std; const int inf=1e9; struct edge{int t,v;edge*next;}e[MAXN<<1],*h[MAXN],*o=e; void add(int x,int y,int vul){o->t=y;o->v=vul;o->next=h[x];h[x]=o++;} ll read(){ ll x=0,f=1;char ch=getchar(); while(!isdigit(ch)){if(ch=='-')f=-1;ch=getchar();} while(isdigit(ch))x=x*10+ch-'0',ch=getchar(); return x*f; } int n,key,rt,base; int sz[MAXN],maxx[MAXN]; int dp1[MAXN<<1],dp2[MAXN<<1],dis[MAXN],Vis[MAXN<<2]; bool vis[MAXN]; ll ans;
void get_root(int x,int pre){ sz[x]=1;maxx[x]=0; link(x){ if(j->t==pre||vis[j->t])continue; get_root(j->t,x); sz[x]+=sz[j->t];maxx[x]=max(maxx[x],sz[j->t]); } maxx[x]=max(maxx[x],base-sz[x]); if(key>maxx[x])key=maxx[x],rt=x; }
int tot,St[MAXN],tot1; pii st[MAXN];
void dfs(int x,int pre){ St[++tot1]=dis[x]; if(Vis[dis[x]+n])ans+=dp1[n-dis[x]],ans+=dp2[n-dis[x]],st[++tot]=mp(dis[x],1); else ans+=dp2[n-dis[x]],st[++tot]=mp(dis[x],2); if(!dis[x]&&Vis[n])ans++; if(!dis[x]&&!Vis[n])ans+=dp1[n]; Vis[dis[x]+n]++; link(x){ if(j->t==pre||vis[j->t])continue; dis[j->t]=dis[x]+(j->v==0?1:-1); dfs(j->t,x); } Vis[dis[x]+n]--; }
void solve(int x){ vis[x]=1;tot1=0; link(x){ if(vis[j->t])continue; tot=0;dis[j->t]=(j->v==0?1:-1);dfs(j->t,x); inc(i,1,tot)if(st[i].second==1)dp2[n+st[i].first]++;else dp1[n+st[i].first]++; } inc(i,1,tot1)dp1[n+St[i]]=dp2[n+St[i]]=0; link(x){ if(vis[j->t])continue; key=inf;base=sz[j->t];get_root(j->t,0); solve(rt); } }
int main(){ n=read(); int x,y,z; ans=0; inc(i,2,n)x=read(),y=read(),z=read(),add(x,y,z),add(y,x,z); key=inf;base=n;get_root(1,0); solve(rt); printf("%lld\n",ans); return 0; }
|