题意:给出一棵树,每条边的权值要么是0, 要么是1;
若有序数对<u, v>, u到v的路径中满足下列任一条件, 就称该有序数对有效:
- 所经过的路径权值全为0;
- 所经过的路径权值全为1;
- 经过权值为1的路径后不会经过权值为0的路径,即在经过的路径中前半段权值全为1,后半段权值全为0;
问:有多少呢有效有序数对;(注: <a, b> != <b, a>)
定义一个dp[n][4];
dp[u][0]表示以u为根的树中到子节点到u的路径全为0的点的个数;
dp[u][1]表示以u为根的树中到子节点到u的路径前半程为1, 后半程为0的点的个数;
dp[u][2]表示以u为根的树中到子节点到u的路径前半程为0, 后半程为1的点的个数;
dp[u][3]表示以u为根的树中到子节点到u的路径全为1的点的个数;
再来个num[4];
num[0]表示到父节点的路径全为0的点的个数;
num[1]表示到父节点的路径前半程为1, 后半程为0的点的个数;
num[2]表示到父节点的路径前半程为0, 后半程为1的点的个数;
num[3]表示到父节点的路径全为1的点的个数;
对于点u,在已知以上条件的情况下,只需拼接路径就可以了;
dp[u][0]与num[0], num[2], num[3]拼接;
dp[u][1]与num[3]拼接;
dp[u][2]与num[0]拼接;
dp[u][3]与num[0], num[1], num[3]拼接;
#include <bits/stdc++.h>
using namespace std;
const int maxn=2e5+10;
struct node{
int v, w, nxt;
}edge[maxn*2];
int head[maxn];
int cnt;
void add(int u, int v, int w){
edge[cnt].v=v;
edge[cnt].w=w;
edge[cnt].nxt=head[u];
head[u]=cnt++;
}
long long dp[maxn][5];
long long num[5];
long long ans;
void dfs(int u, int fa){
//cout << u << ' ' << fa << endl;
for(int i=head[u]; i!=-1; i=edge[i].nxt){
int v=edge[i].v;
int w=edge[i].w;
if(v==fa) continue;
dfs(v, u);
if(w==0){
num[0]=dp[v][0]+1;
num[1]=0;
num[2]=dp[v][2]+dp[v][3];
num[3]=0;
}
else{
num[0]=0;
num[1]=dp[v][0]+dp[v][1];
num[2]=0;
num[3]=dp[v][3]+1;
}
ans+=dp[u][0]*(num[0]*2+num[2]+num[3]);
ans+=dp[u][1]*num[3];
ans+=dp[u][3]*(num[3]*2+num[1]+num[0]);
ans+=dp[u][2]*num[0];
for(int i=0; i<4; i++)
dp[u][i]+=num[i];
}
ans+=dp[u][0]*2+dp[u][3]*2+dp[u][1]+dp[u][2];
}
int main(){
ios::sync_with_stdio(false);
int n;
cin >> n;
memset(head, -1, sizeof(head));
cnt=0;
for(int i=1; i<n; i++){
int u, v, w;
cin >> u >> v >> w;
add(u, v, w);
add(v, u, w);
}
ans=0;
memset(dp, 0, sizeof(dp));
dfs(1, 0);
cout << ans << endl;
return 0;
}