1156D - 0-1-Tree 树形DP

题意：给出一棵树，每条边的权值要么是0， 要么是1；
若有序数对<u, v>, u到v的路径中满足下列任一条件， 就称该有序数对有效：

1. 所经过的路径权值全为0；
2. 所经过的路径权值全为1；
3. 经过权值为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;
}