Codeforces Gym102448J Jingle Bells

题目链接

Codeforces Gym102448J Jingle Bells

题目大意

给定一棵有$n$个节点的树，树的边一共有$5$种颜色，刚开始树的有些边已经被染色了。现在要你把剩下没染色的边染色，让两条具有公共节点的相邻的边的颜色不同，有多少种方案，对$10^9+7$取模。

解题思路

树形$dp$，可以看出这题要在树上计数，$dp$。

发现树的边一共有$5$种颜色，非常少，那么我们可以状压。 我们先转化一下，将“两条相邻的边的颜色不同”转化成“任意一个点连出的边的颜色都不相同”。于是我们设$f_{x,sta}$表示以$x$为子树的根，从$x$连出的边的颜色状态为$sta$时，染色的方案数。

然后我们考虑状态转移方程。设$so{n}_x$为$x$的子节点集合，$y$为$x$的当前计算的子节点，$pos$为$y$节点的边染色状态，当前新加的边颜色为$color$，则
$$f_{x,sta\mid 2^{color}}=\sum _{y\in so{n}_x}{f}_{x,sta}\times f_{y,pos(pos\&2^{color}=0)}$$
表示为新推出的状态由现在状态的方案数乘子节点不包含新增边颜色的所有状态的方案数。

注意点

树形$dp$顺序

从大到小，因为从小到大会重复计算。在$sta$状态的时候，往$sta\mid 2^{color}$状态加入了方案数，而$sta\mid 2^{color}>sta$，所以往后转移的时候方案数会重复叠加产生错误。

清零

在用$f_{x,sta}$转移$f_{x,sta|2^{color}}$后，将$f_{x,sta}$赋值为$0$。 这是为什么呢？ 因为$x$可能有多个子节点$y$，由$f_{x,sta}$和$f_{y,pos}$转移后，$f_{x,sta|2^{color}}$就被更新了，此时的$f_{x,sta}$不能继续使用，因为枚举顺序是从大到小，导致了此时使用的$f_{x,sta}$实际上是上一个子节点$y_1$留下的方案数，和$01$背包的从大到小的思想类似，相当于把$f_{x,sta}$里的每种方案数都$\times f_{y,pos}$，乘法原理。上一步是$y_1$子节点及之前的答案，这一步新增了$y$节点，那么就是$y$节点的方案$\times y_1$节点及之前的所有方案。

更好理解的方法，假设节点$x$有$3$个子节点$y_1,y_2,y_3$，我们已经出来出$y_1$更新的方案数了，这时假设我们还不知道有$y_2,y_3$，新开辟了一个$y_2$节点，那么我们原来的方案数就要乘$y2$和$x$这条边的染色方案数，然后原来的方案数是没$y_2$的时候的，用完后清零，再由$f_{x,st{a}_1(st{a}_1\in sta)}$后面状态的方案数去更新本次新增节点$y_2$后的方案数。

代码实现

#include<bits/stdc++.h>
#define For(i,a,b)for(int i=a;i<=b;i++)
using namespace std;
typedef long long ll;
const int N=2e5+5,mod=1e9+7,up=31;
ll f[N][32],ans;
int n,cnt,x,y,z,a[N],b[N],c[N],d[N];
void add(int x,int y,int z){a[++cnt]=y,b[cnt]=c[x],c[x]=cnt,d[cnt]=z;}
void dfs(int x,int fa){
    f[x][0]=1;//什么颜色都不染方案数为1
    for(int i=c[x];i;i=b[i])if(a[i]!=fa){
        dfs(a[i],x);
        int y=a[i],l=0,r=4;//五种颜色
        if(d[i])l=r=d[i]-1;//假如这条边本身有色就设为自己
        for(int sta=up;sta>=0;f[x][sta--]=0)//倒序枚举状态
            for(int c=l;c<=r;c++)if(!(f[x][sta]==0||(sta&(1<<c))))//这条边颜色
                for(int son=0;son<=up;son++)if(!(son&(1<<c)))//子节点状态
                    f[x][sta|(1<<c)]=(f[x][sta|(1<<c)]+f[x][sta]*f[y][son]%mod)%mod;//增加方案
    }
}
int main(){
    scanf("%d",&n);
    For(i,2,n)scanf("%d%d%d",&x,&y,&z),add(x,y,z),add(y,x,z);
    dfs(1,0);For(i,0,up)(ans+=f[1][i])%=mod;//f[1]中为答案每种状态都行，无限制
    return printf("%lld\n",ans),0;
}