codeforces1156D 0-1-Tree 换根dp

题目传送门

题意：

　　给定一棵n个点的边权为0或1的树，一条合法的路径(x,y)(x≠y)满足，从x走到y，一旦经过边权为1的边，就不能再经过边权为0的边，求有多少边满足条件？

思路：

　　首先，这道题也可以用并查集的做法过，点这里

　　那换根dp怎么写呢？

　　设$f[u]$为以1为根，自下而上到$u$的末节点是1的合法路径数量，$g[u]$代表以1为根，自下而上到$v$末节点是0的合法路径数量，这个可以通过一遍dfs简单求解。

　　再设$nf[u]$和$ng[u]$代表以u为根的两种合法路径数量，进行换根dfs，在换根的过程中：

　　若某一条边是0边，则：

　　　　$ng[st.to]=ng[u]$，$nf[st.to]=f[st.to]$。这个方程也很好理解，白边的路径是不会变的，所有从父节点自上而下转移过来的黑边到了这里都是非法路径了。

　　若某一条边是1边，则：

　　　　$ng[st.to]=g[st.to]$,$nf[st.to]=nf[u]-g[st.to]+ng[u]$,白边只有从下往上过来的了。黑边要减去  到当前位置为白边与父节点的黑边连接形成的边 ，再加上父节点是白边，加上黑边形成的边。

#pragma GCC optimize (2)
#pragma G++ optimize (2)
#pragma comment(linker, "/STACK:102400000,102400000")
#include<bits/stdc++.h>
#include<cstdio>
#include<vector>
#define rep(i,a,b) for(int i=a;i<=b;i++)
#define dep(i,b,a) for(int i=b;i>=a;i--)
#define clr(a,b) memset(a,b,sizeof(a))
#define pb push_back
#define pii pair<int,int >
using namespace std;
typedef long long ll;
const int maxn=200010;
ll rd()
{
    ll x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
int T;
struct edge{
    int to,w;
};
vector<edge>ve[maxn];
int f[maxn],g[maxn],nf[maxn],ng[maxn];
int n,m;
ll ans;
void dfs_1(int u,int fa){
    for(auto &st:ve[u]){
        if(st.to==fa)continue;
        dfs_1(st.to,u);
        if(st.w==0){
            g[u]+=g[st.to]+1;
        }else{
            f[u]+=f[st.to]+1+g[st.to];
        }
    }
}
void dfs_2(int u,int fa){
    ans+=nf[u]+ng[u];
    for(auto &st:ve[u]){
        if(st.to==fa)continue;
        if(st.w==0){
            ng[st.to]=ng[u];
            nf[st.to]=f[st.to];
        }else{
            ng[st.to]=g[st.to];
            nf[st.to]=nf[u]-g[st.to]+ng[u];
        }
        dfs_2(st.to,u);
    }
}
int main(){
    cin>>n;
    rep(i,1,n-1){
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        ve[u].pb({v,w});
        ve[v].pb({u,w});
    }
    dfs_1(1,0);
    nf[1]=f[1],ng[1]=g[1];
    dfs_2(1,0);
    cout<<ans<<endl;
}