Codeforces D. 0-1-Tree (树形dp)

https://codeforces.com/contest/1156/problem/D

题意：给你一颗树，同时给每一条边定义一个属性0和1，定义（x，y）表示x到y，其中我们经过1边后不能经过0边，求有多少对（x，y）；

做法：本蒟蒻只想得到树形dp的想法，想不到大佬的并查集方法。dp[u][4]，分别定义了u和他子树的四种状态。如下图

dp[u][0]表示u的子树中全部是0连在一起的个数。

dp[u][1]表示u的子树中全部的与u相连上面的一部分是1，下面的一部分是0

dp[u][2]表示u的子树中全部的与u相连上面的一部分是0，下面的一部分是1.

dp[u][3]表示u的子树上面全是1相连的

明白了这个就不难了，代码虽然看起来有一点点长，但不过很好理解。代码怎么写就各看各的本事了

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=200010;
int n;
vector<pair<int,int> >g[N];
ll ans=0,dp[N][4];
void dfs(int u,int fa)
{
    ll cnt[4]={0};
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i].first;
        int w=g[u][i].second;
        if(v==fa)
            continue;
        dfs(v,u);
        memset(cnt,0,sizeof(cnt));
        if(w==0)
        {
            cnt[0]=dp[v][0]+1;cnt[1]=0;
            cnt[2]=dp[v][2]+dp[v][3];
        }
        else
        {
            cnt[1]=dp[v][0]+dp[v][1];
            cnt[3]=dp[v][3]+1;cnt[2]=0;
        }
        ans+=dp[u][0]*(2*cnt[0]+cnt[2]+cnt[3]);
        ans+=dp[u][2]*cnt[0];
        ans+=dp[u][3]*(cnt[0]+cnt[1]+2*cnt[3]);
        ans+=dp[u][1]*cnt[3];
        for(int i=0;i<4;i++)
            dp[u][i]+=cnt[i];
    }
    ans+=(dp[u][0]+dp[u][3])*2+dp[u][1]+dp[u][2];
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    cin>>n;
    int u,v,w;
    for(int i=1;i<n;i++)
    {
        cin>>u>>v>>w;
        g[u].push_back(make_pair(v,w));
        g[v].push_back(make_pair(u,w));
    }
    memset(dp,0,sizeof(dp));
    dfs(1,0);
    cout<<ans<<endl;
    return 0;
}