【CF1010F】Tree(链分治)(分治NTT)

传送门


时限7s的题的复杂度令人捉摸不透。。。

O ( n log ⁡ 3 n ) O(n\log^3n) O(nlog3n)拿下CFrk1(仗着这道题分治FFT部分常数小),好像std也是 O ( n log ⁡ 3 n ) O(n\log^3 n) O(nlog3n)的。。。

问题简述如下:

一颗二叉树,根节点为 1 1 1,允许你砍掉任意子树,保留根节点所在的连通块,然后要求你设置每个点的权值,要求每个点权值不小于它的两个儿子的权值之和,且根节点权值为 x x x的方案数,连通块不一样或者点的权值不一样都算不同方案。

题解:

首先第一步转化,题目要求每个点的权值不小于儿子的权值之和,转换一下每个点减去两个儿子的权值,实际上是这样一个问题:给每个点赋一个非负权值,且权值和为 x x x(即原题中根节点的限制),发现这个和树的形态没有关系只和树的siz有关系,插板法可以算。

那么现在问题就是树上选一个包含根节点的联通块,大小为 k k k的方案数。

树形背包可以 O ( n 2 ) O(n^2) O(n2)算,但是注意到可以转化成生成函数的形式,设 F u ( x ) F_u(x) Fu(x)表示以 u u u为根,选择大小为 k k k的联通块的方案数。

轻重链剖分之后对于每条重链分治NTT算出链顶的答案即可。


代码:

#include
#define ll long long
#define re register
#define cs const

using std::cerr;
using std::cout;

cs int mod=998244353;
inline int add(int a,int b){a+=b-mod;return a+(a>>31&mod);}
inline int dec(int a,int b){a-=b;return a+(a>>31&mod);}
inline int mul(int a,int b){ll r=(ll)a*b;return r>=mod?r%mod:r;}
inline int power(int a,int b,int res=1){
	for(;b;b>>=1,a=mul(a,a))(b&1)&&(res=mul(res,a));
	return res;
}
inline void Inc(int &a,int b){a+=b-mod;a+=a>>31&mod;}
inline void Dec(int &a,int b){a-=b;a+=a>>31&mod;}
inline void Mul(int &a,int b){a=mul(a,b);}

typedef std::vector<int> Poly;

cs int bit=19,SIZE=1<<bit|1;

int r[SIZE],*w[bit+1],inv[SIZE];
inline void init_NTT(){
	for(int re i=1;i<=bit;++i)w[i]=new int[1<<i-1];
	int wn=power(3,mod-1>>bit);w[bit][0]=1;
	for(int re i=1;i<(1<<bit-1);++i)w[bit][i]=mul(w[bit][i-1],wn);
	for(int re i=bit-1;i;--i)
	for(int re j=0;j<(1<<i-1);++j)w[i][j]=w[i+1][j<<1];
	inv[0]=inv[1]=1;
	for(int re i=2;i<SIZE;++i)inv[i]=mul(mod-mod/i,inv[mod%i]);
}
inline void NTT(int *A,int len,int typ){
	for(int re i=1;i<len;++i)if(i<r[i])std::swap(A[i],A[r[i]]);
	for(int re i=1,d=1;i<len;i<<=1,++d)
	for(int re j=0;j<len;j+=i<<1)
	for(int re k=0;k<i;++k){
		int &t1=A[j+k],&t2=A[i+j+k],t=mul(t2,w[d][k]);
		t2=dec(t1,t);Inc(t1,t);
	}
	if(typ==-1){
		std::reverse(A+1,A+len);
		for(int re i=0,iv=inv[len];i<len;++i)Mul(A[i],iv);	
	}
}
inline void NTT(Poly &A,int len,int typ){NTT(&A[0],len,typ);}
inline void init_rev(int l){for(int re i=1;i<l;++i)r[i]=r[i>>1]>>1|((i&1)?l>>1:0);}

inline Poly operator*(Poly a,Poly b){
	int n=a.size(),m=b.size(),deg=n+m-1,l=1;
	while(l<deg)l<<=1;init_rev(l);
	a.resize(l),NTT(a,l,1);
	b.resize(l),NTT(b,l,1);
	for(int re i=0;i<l;++i)Mul(a[i],b[i]);
	NTT(a,l,-1);return a.resize(deg),a;
}

inline Poly operator+(cs Poly &a,cs Poly &b){
	Poly c=a;if(a.size()<b.size())c.resize(b.size());
	for(int re i=0,li=b.size();i<li;++i)Inc(c[i],b[i]);
	return c;
}

cs int N=1e5+7;
int n;ll m;
std::vector<int> G[N];

int fa[N],siz[N],son[N];

Poly a[N],p[N];

void dfs(int u,int p){
	siz[u]=1,fa[u]=p;
	for(int re v:G[u])if(v!=p){
		dfs(v,u);siz[u]+=siz[v];
		if(siz[v]>siz[son[u]])son[u]=v;
	}
}

inline void solve(int l,int r,Poly &f,Poly &g){
	if(l==r){f=g=a[l];return ;}
	int mid=l+r>>1;Poly fl,fr,gl,gr;
	solve(l,mid,fl,gl);solve(mid+1,r,fr,gr);
	f=fl*fr,g=fl*gr+gl;
}

Poly dfs2(int u){
	for(int w=u;w;w=son[w]){
		p[w].clear();for(int re v:G[w])
		if(v!=fa[w]&&v!=son[w])p[w]=dfs2(v);
		if(p[w].empty())p[w].push_back(0);
		p[w][0]++,p[w].insert(p[w].begin(),0);
	}
	int k=0;for(int re v=u;v;v=son[v])a[++k].swap(p[v]);
	Poly f,g;solve(1,k,f,g);return g;
}

signed main(){
#ifdef zxyoi
	freopen("tree.in","r",stdin);
#endif
	init_NTT();
	scanf("%d%lld",&n,&m);m%=mod;
	for(int re i=1;i<n;++i){
		int u,v;scanf("%d%d",&u,&v);
		G[u].push_back(v);
		G[v].push_back(u);
	}dfs(1,0);Poly res=dfs2(1);
	int C=1,ans=0;
	for(int re i=1,li=res.size();i<li&&C;++i)
	Inc(ans,mul(C,res[i])),Mul(C,mul(add(m,i),inv[i]));
	cout<<ans<<"\n";
	return 0;
}

你可能感兴趣的:(树分治,生成函数)