You are given a tree (an undirected connected acyclic graph) consisting of n vertices and n−1 edges. A number is written on each edge, each number is either 0 (let’s call such edges 0-edges) or 1 (those are 1-edges).
Let’s call an ordered pair of vertices (x,y) (x≠y) valid if, while traversing the simple path from x to y, we never go through a 0-edge after going through a 1-edge. Your task is to calculate the number of valid pairs in the tree.
Input
The first line contains one integer n (2≤n≤200000) — the number of vertices in the tree.
Then n−1 lines follow, each denoting an edge of the tree. Each edge is represented by three integers xi, yi and ci (1≤xi,yi≤n, 0≤ci≤1, xi≠yi) — the vertices connected by this edge and the number written on it, respectively.
It is guaranteed that the given edges form a tree.
Output
Print one integer — the number of valid pairs of vertices.
Example
Input
7
2 1 1
3 2 0
4 2 1
5 2 0
6 7 1
7 2 1
Output
34
思路:
显然树形dp,一开始我以为只是求链都是黑色的路径的数量,发现这个怎么这么简单,发现读错了题,题目实际上求得是是黑白两条拼接的路径数,就是一侧都是白色,另一侧都是黑色,也可以全黑或者全白。
然后我设两个状态,就是全黑和全白,发现貌似还少了一些case,然后我就又设了一个上头是黑色,下面至少有一头白色。然后这个我在实际去写的时候,当在白边转移时,自然而想到少了另一个状态,就是上面是白色,下面至少有一头黑色。
回过头来看,其实发现,白色和黑色的地位是等价的,根据对称性,就应该能设计出偶数个状态,所有最终是4个状态,在转移时,遇到黑色和遇到白色,的转移方程也是对称的。
然后智商下线了,忘记把全黑和全白的情况乘二了。。。怎么算都还是少了。。。自己手动枚举样例发现和我做也是一样。。。对自己无语了。。。注意是ordered pair
状态:
f
[
u
]
[
0
]
f[u][0]
f[u][0]为以u为子树的时候,u为端点,边全是白色的pair个数
f
[
u
]
[
1
]
f[u][1]
f[u][1]为以u为子树的时候,u为端点,边全是黑色的pair个数
f
[
u
]
[
2
]
f[u][2]
f[u][2]为以u为子树的时候,u为端点,靠近u一头边为白色,终点是白色的pair个数
f
[
u
]
[
3
]
f[u][3]
f[u][3]为以u为子树的时候,u为端点,靠近u一头边为黑色,终点是黑色的pair个数
代码:
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<iostream>
#include<map>
#include<cmath>
#include<vector>
#include<set>
#define INF 0x7f7f7f7f
#define maxx 200005
using namespace std;
typedef long long ll;
int head[maxx],to[maxx<<1],_next[maxx<<1],w[maxx<<1];
int edge;
void addEdge(int x,int y,int c)
{
to[++edge]=y,w[edge]=c,_next[edge]=head[x],head[x]=edge;
to[++edge]=x,w[edge]=c,_next[edge]=head[y],head[y]=edge;
}
int f[maxx][4];
int n;
ll ans=0;
void dfs(int u,int fa)
{
for(int i=head[u];i;i=_next[i])
{
int v=to[i];
if(v==fa)continue;
dfs(v,u);
if(w[i])
{
ans+=(ll)2*f[u][1]*(f[v][1]+1);
ans+=(ll)(f[u][0]+f[u][3])*(f[v][1]+1);
ans+=(ll)f[u][1]*(f[v][0]+f[v][3]);
f[u][1]+=f[v][1]+1;
f[u][3]+=f[v][3];
f[u][3]+=f[v][0];
}
else
{
ans+=(ll)2*f[u][0]*(f[v][0]+1);
ans+=(ll)(f[u][1]+f[u][2])*(f[v][0]+1);
ans+=(ll)f[u][0]*(f[v][1]+f[v][2]);
f[u][0]+=f[v][0]+1;
f[u][2]+=f[v][2];
f[u][2]+=f[v][1];
}
}
ans+=2*f[u][0];
ans+=2*f[u][1];
ans+=f[u][2];
ans+=f[u][3];
}
int main()
{
cin>>n;
int x,y,c;
for(int i=1;i<n;i++)
{
scanf("%d%d%d",&x,&y,&c);
addEdge(x,y,c);
}
dfs(1,0);
cout<<ans<<endl;
return 0;
}