UFPE Starters Final Try-Outs 2020 J.Jingle Bells 树形dp

题目链接 http://codeforces.com/gym/102448/problem/J

题意

你现在有一棵树，5种颜色 $(1,2,3,4,5)$ ，树的边会有一种颜色或者没有被染色。现在要你给树染上颜色，让每一个顶点的所有边都带有不同的颜色，问你有多少中方案。

做法

树形 $dp$ ，（赛上并没有时间想，挺神奇的一道我应该是做不出来的一道题）。 $dp[u][s]$ 表示的是从顶点开始遍历从下往上遍历到点 $u$ 的时候，目前其边的颜色状态为 $s$ 的方案数，然后在遍历的过程中，要枚举当前边的颜色，来将这条边的两个端点的状态进行相乘。

在遍历的过程中，注意遍历的顺序，是从大到小的，因为在更新的过程中，是将答案加至 $|$ 后的值上的，所以在不仅要从大到小，并且要将当前的值清空，以供接下来的值来进行增加。

最后所有的答案都保存在先遍历的点 $1$ 中，将答案加起来即可。

代码

#include <bits/stdc++.h>
#define rep(i,a,b) for(int i=a;i<=b;i++)
using namespace std;
typedef long long ll;
const int maxn=100005;
const int maxm=maxn*2;
const int mod=(int)1e9+7;
const int up=(1<<5)-1;
ll dp[maxn][40];
int nex[maxm],to[maxm],head[maxn];
int cnt,col[maxm],n;
void add(int u,int v,int c){
    to[cnt]=v,nex[cnt]=head[u];
    col[cnt]=c,head[u]=cnt++;
}
void dfs(int u,int f){
    dp[u][0]=1;
    for(int i=head[u];~i;i=nex[i]){
        int v=to[i];
        if(v==f) continue;
        dfs(v,u);
        int l=0,r=4;
        if(col[i]) l=r=col[i]-1;
        for(int fa=up;fa>=0;fa--){
            for(int co=l;co<=r;co++){

                if(dp[u][fa]==0||(fa&(1<<co))) continue;
                for(int son=0;son<=up;son++){

                    if(son&(1<<co)) continue;
                    dp[u][fa|(1<<co)]=(dp[u][fa|(1<<co)]+dp[u][fa]*dp[v][son]%mod)%mod;
                }

            }
            dp[u][fa]=0;
        }
    }
}
int main(){
    memset(head,-1,sizeof(head));
    scanf("%d",&n);
    rep(i,2,n){
        int x,y,z; scanf("%d%d%d",&x,&y,&z);
        add(x,y,z); add(y,x,z);
    }
    dfs(1,-1);
    ll ans=0;
    rep(i,0,up){
        ans=(ans+dp[1][i])%mod;
    }
    printf("%lld\n",ans);
    return 0;
}