【PKUSC2019】树染色【线段树合并】【树形DP】

左手的ㄟ右手 2021-12-23 12:41 266阅读 0赞

## Description ##

给出一棵n个点的树，现在有m种颜色，要给每个节点染色，相邻节点不能同色。  
另外有k条限制，形如x号点不能为颜色y  
同一节点有可能有多条限制。  
求方案数对998244353取模的结果。

n<=200000,m<=1e9,k<=400000

## Solution ##

考场上一直在想怎么容斥做，怎么都弄不出来。  
学傻了。

考虑暴力DP  
设\\(f\[i\]\[j\]\\)为当前处理了以i为根的子树，i的颜色为j的方案数。  
记\\(g\[i\]=\\sum\\limits\_\{k\}f\[i\]\[k\]\\)

显然有转移\\\[f\[i\]\[j\]=\[!ban\[i\]\[j\]\]\\prod\_\{p\\in son\[i\]\}(g\[p\]-f\[p\]\[j\])\\\]

但是这样的状态数是\\(n\*m\\)的，我们发现只需要记下子树中有的颜色，其他的颜色的答案都是一样的。  
这样状态数缩减到\\(n\*k\\)，但还是很大，于是我们考虑采用线段树来维护。

转移的时候我们将子树一个个的合并到根  
大概是\\(f\[i\]\[j\]=(g\[p\]-f\[p\]\[j\])\*f\[i\]\[j\]\\)  
根据这个我们就可以线段树合并了。  
如果只有父亲有，就直接乘  
儿子父亲都有暴力合并  
只有儿子有的话把括号拆开，就是乘上\\(-f\[i\]\[j\]\\)加上\\(g\[p\]\*f\[i\]\[j\]\\)  
需要维护区间乘区间加，类似一次函数维护即可。  
时间复杂度大概是\\(O((n+k)\\log m)\\)，具体可以看代码。

## Code ##

写了个对拍没问题，姑且当它是对的吧

#include <bits/stdc++.h> #define fo(i,a,b) for(int i=a;i<=b;++i) #define fod(i,a,b) for(inti i=a;i>=b;--i) #define N 200005 #define M 13000005 #define LL long long #define mo 998244353 using namespace std; int n,m,l,fs[N],nt[2*N],dt[2*N],m1; vector<int> qs[N]; LL ksm(LL k,LL n) { LL s=1; for(;n;n>>=1,k=k*k%mo) if(n&1) s=s*k%mo; return s; } int n1,t[M][2],sz[M],rt[N]; LL sp[M],g[N],f[N],lz[M][2]; void nwp(int &k) { if(!k) k=++n1,lz[k][0]=1,lz[k][1]=0; } void ins(int k,int l,int r,int x,int v) { if(l==r) {sp[k]=0,sz[k]=1;return;} int mid=(l+r)>>1; if(x<=mid) nwp(t[k][0]),ins(t[k][0],l,mid,x,v); else nwp(t[k][1]),ins(t[k][1],mid+1,r,x,v); sp[k]=(sp[t[k][0]]+sp[t[k][1]]); if(sp[k]>=mo) sp[k]-=mo; sz[k]=sz[t[k][0]]+sz[t[k][1]]; } LL gp,fp,fk,vs; void upd(int k,LL u,LL v) { sp[k]=(u*sp[k]+v*sz[k])%mo; lz[k][0]=lz[k][0]*u%mo; lz[k][1]=(lz[k][1]*u%mo+v)%mo; } void down(int k) { if(lz[k][0]!=1||lz[k][1]!=0) { if(t[k][0]) upd(t[k][0],lz[k][0],lz[k][1]); if(t[k][1]) upd(t[k][1],lz[k][0],lz[k][1]); lz[k][0]=1,lz[k][1]=0; } } void mrg(int &k,int x,int l,int r) { if(!k) { if(!x) return; k=x,upd(k,mo-fk,gp*fk%mo); return; } if(!x) {upd(k,(gp-fp+mo)%mo,0);return;} if(l==r) {sp[k]=(gp-sp[x]+mo)%mo*sp[k]%mo,sz[k]=sz[k]|sz[x];return;} int mid=(l+r)>>1; down(k),down(x); mrg(t[k][0],t[x][0],l,mid); mrg(t[k][1],t[x][1],mid+1,r); sp[k]=(sp[t[k][0]]+sp[t[k][1]])%mo; sz[k]=sz[t[k][0]]+sz[t[k][1]]; } void dfs(int k,int fa) { f[k]=1; nwp(rt[k]); int r=qs[k].size(); fo(j,0,r-1) ins(rt[k],1,m,qs[k][j],0); for(int i=fs[k];i;i=nt[i]) { int p=dt[i]; if(p!=fa) { dfs(p,k); gp=g[p],fk=f[k],fp=f[p]; mrg(rt[k],rt[p],1,m); f[k]=(g[p]-f[p]+mo)%mo*f[k]%mo; } } g[k]=(f[k]*(LL)(m-sz[rt[k]])%mo+sp[rt[k]])%mo; } void link(int x,int y) { nt[++m1]=fs[x]; dt[fs[x]=m1]=y; } int main() { cin>>n>>m>>l; fo(i,1,n-1) { int x,y; scanf("%d%d",&x,&y); link(x,y),link(y,x); } fo(i,1,l) { int x,y; scanf("%d%d",&x,&y); qs[x].push_back(y); } dfs(1,0); printf("%lld\n",g[1]); }

转载于:https://www.cnblogs.com/BAJimH/p/10942994.html