CodeForces - 1156D 0-1-Tree

题意：给出一棵树，树的边权只有0和1。求有多少有序点对，其最短路径上每条权值为0的边不紧跟在权值为1的边后面。

解：合法路径如下所示：

000000 
111111 
000111 

随便找个结点为根节点，形成一棵树。结点a_i为其中一棵子树T_i的根。考虑统计以结点a_i为终点，最后一条边为0/1的路径数量。T_i的每棵子树内答案可以直接得到，子树间就是把两条路径拼起来。因此再添加一个不合法状态111000用来转移。给几个状态编号（从左往右为树中从下到上）：

考虑要不要换根的问题。如果将a_i作为整棵树的根，那它的原来的父节点就会成为它的一棵子树的根。这颗子树也应该和原来的子树拼一拼。但这种情况会在处理祖先节点的时候覆盖掉，所以dfs一遍，考虑目前它几棵子树的关系就可以了。

接下来拼几棵子树的路径。假设一棵子树和a_i之间边w_i的边权为1，那么这颗子树可以是状态0、1和2；状态0加上w_i变成状态2，能和状态1拼；状态1能和状态1、状态0、状态2拼，状态2能和状态1拼。边权为1同理。

代码：

#include <bits/stdc++.h>
using namespace std;
#define maxx 200005
#define maxn 25
#define maxm 205
#define ll long long
#define inf 1000000009
#define mod 2520
struct edge{
    int u,v,w,next;
}a[maxx*2];
int head[maxx]={0},cnt=0;
void add(int u,int v,int w){
    a[++cnt].u=u;a[cnt].v=v;a[cnt].w=w;a[cnt].next=head[u];head[u]=cnt;
}

ll ans=0;
int n;
ll dp[maxx][4]={0};
//000000  0
//111111  1
//000111  2
//111000  3
void dfs1(int now,int fa){
    for(int i=head[now];i;i=a[i].next){
        int to=a[i].v;
        if(to==fa) continue;
        dfs1(to,now);
        if(a[i].w==1){
            ans+=dp[now][1]*(dp[to][1]+1)*2;
            ans+=dp[now][1]*dp[to][2];
            ans+=dp[now][1]*dp[to][0];
            ans+=dp[now][0]*(dp[to][1]+1);
            ans+=dp[now][2]*(dp[to][1]+1);

            dp[now][1]+=dp[to][1]+1;
            dp[now][2]+=dp[to][2];
            dp[now][2]+=dp[to][0];
        }
        else if(a[i].w==0){
            ans+=dp[now][0]*(dp[to][0]+1)*2;
            ans+=dp[now][0]*dp[to][3];
            ans+=dp[now][0]*dp[to][1];
            ans+=dp[now][1]*(dp[to][0]+1);
            ans+=dp[now][3]*(dp[to][0]+1);

            dp[now][0]+=dp[to][0]+1;
            dp[now][3]+=dp[to][1];
            dp[now][3]+=dp[to][3];
        }
    }
    ans+=dp[now][0]*2;
    ans+=dp[now][1]*2;
    ans+=dp[now][2];
    ans+=dp[now][3];
}
signed main() {
    scanf("%d",&n);
    for(int i=1;i<=n-1;i++){
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);
        add(y,x,z);
    }
    dfs1(1,0);
    printf("%lld\n",ans);
    return 0;
}

还有一种做法是统计一个结点a_i能从边权为0/1的边出发，合法地到达多少个点。先从下到上dfs一遍，然后从上到下统计答案。

反正算起来都很绕啦