[Codechef]-Walk on Tree (BM+constant coefficient linear recurrence)

Posted by sarathi on Mon, 07 Oct 2019 17:18:12 +0200

Portal

The practice of O(n3logk)O(n^3logk)O(n3logk) is obvious.

Consider that it is actually a requirement to sum some positions of a matrix.
In fact, this one has a recursive formula, and the recursive formula is the characteristic polynomial of the matrix.

In fact, it is very obvious that every time the modulus of the characteristic polynomial is taken.
One location is the usual constant coefficient linear recurrence.
The same is true for the sum of some positions.

Violent DPDP pulls out the first nnn item and BMBMBM BM water passes by

#include<bits/stdc++.h>
using namespace std;
#define gc getchar
inline int read(){
	char ch=gc();
	int res=0,f=1;
	while(!isdigit(ch))f^=ch=='-',ch=gc();
	while(isdigit(ch))res=(res+(res<<2)<<1)+(ch^48),ch=gc();
	return f?res:-res;
}
#define pb push_back
#define cs const
#define ll long long
#define poly vector<int>
#define bg begin
cs int mod=998244353,G=3;
inline int add(int a,int b){return (a+=b)>=mod?a-mod:a;}
inline void Add(int &a,int b){(a+=b)>=mod?(a-=mod):0;}
inline int dec(int a,int b){return (a-=b)<0?a+mod:a;}
inline void Dec(int &a,int b){(a-=b)<0?(a+=mod):0;}
inline int mul(int a,int b){return 1ll*a*b>=mod?1ll*a*b%mod:a*b;}
inline void Mul(int &a,int b){a=mul(a,b);}
inline int ksm(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 chemx(ll &a,ll b){a<b?a=b:0;}
inline void chemn(ll &a,ll b){a>b?a=b:0;}
cs int N=(1<<17)|5,C=17;
poly w[C+1];
inline void init_w(){
	for(int i=1;i<=C;i++)w[i].resize(1<<(i-1));
	int wn=ksm(G,(mod-1)/(1<<C));
	w[C][0]=1;
	for(int i=1;i<(1<<(C-1));i++)w[C][i]=mul(w[C][i-1],wn);
	for(int i=C-1;i;i--)
	for(int j=0;j<(1<<(i-1));j++)
	w[i][j]=w[i+1][j<<1];
}
int rev[N];
inline void init_rev(int lim){
	for(int i=0;i<lim;i++)rev[i]=(rev[i>>1]>>1)|((i&1)*(lim>>1));
}
inline void ntt(poly &f,int lim,int kd){
	for(int i=0;i<lim;i++)if(i>rev[i])swap(f[i],f[rev[i]]);
	for(int a0,a1,mid=1,l=1;mid<lim;mid<<=1,l++)
	for(int i=0;i<lim;i+=(mid<<1))
	for(int j=0;j<mid;j++)
	a0=f[i+j],a1=mul(w[l][j],f[i+j+mid]),f[i+j]=add(a0,a1),f[i+j+mid]=dec(a0,a1);
	if(kd==-1){
		reverse(f.bg()+1,f.bg()+lim);
		for(int i=0,inv=ksm(lim,mod-2);i<lim;i++)Mul(f[i],inv);
	}
}
inline poly operator +(poly a,poly b){
	int deg=max(a.size(),b.size());
	a.resize(deg),b.resize(deg);
	for(int i=0;i<deg;i++)Add(a[i],b[i]);
	return a;
}
inline poly operator -(poly a,poly b){
	int deg=max(a.size(),b.size());
	a.resize(deg),b.resize(deg);
	for(int i=0;i<deg;i++)Dec(a[i],b[i]);
	return a;
}
inline poly operator *(poly a,poly b){
	int deg=a.size()+b.size()-1,lim=1;
	if(deg<128){
		poly c(deg,0);
		for(int i=0;i<a.size();i++)
		for(int j=0;j<b.size();j++)
		Add(c[i+j],mul(a[i],b[j]));
		return c;
	}
	while(lim<deg)lim<<=1;
	init_rev(lim);
	a.resize(lim),ntt(a,lim,1);
	b.resize(lim),ntt(b,lim,1);
	for(int i=0;i<lim;i++)Mul(a[i],b[i]);
	ntt(a,lim,-1),a.resize(deg);
	return a;
}
inline poly Inv(poly a,int deg){
	poly b(1,ksm(a[0],mod-2)),c;
	for(int lim=4;lim<(deg<<2);lim<<=1){
		c=a,c.resize(lim>>1);
		init_rev(lim);
		b.resize(lim),c.resize(lim);
		ntt(b,lim,1),ntt(c,lim,1);
		for(int i=0;i<lim;i++)Mul(b[i],dec(2,mul(b[i],c[i])));
		ntt(b,lim,-1),b.resize(lim>>1);
	}b.resize(deg);
	return b;
}
inline poly operator /(poly a,poly b){
	int deg=(int)a.size()-(int)b.size()+1;
	if(deg<0)return poly(1,0);
	reverse(a.bg(),a.end());
	reverse(b.bg(),b.end());
	a.resize(deg),b.resize(deg);
	poly c=a*Inv(b,deg);c.resize(deg);
	reverse(c.bg(),c.end());
	return c;
}
inline poly operator %(poly a,poly b){
	if(a.size()<b.size())return a;
	a=a-(a/b)*b;a.resize(b.size()-1);return a;
}
inline poly ksm(poly a,int b,poly res,poly Mod){
	for(;b;b>>=1,a=a*a%Mod)if(b&1)res=res*a%Mod;
	return res;
}
namespace Cas{
	poly f;int n;
	inline void init(poly coef){
		init_w();
		n=coef.size();
		f.resize(n+1);
		for(int i=1;i<=n;i++)f[n-i]=mod-coef[i];
		f[n]=1;
	}
	inline poly calc(int k){
		poly g(2),res(1,1);g[1]=1;
		res=ksm(g,k,res,f);
		return res;
	}
}
namespace B_M{
	poly r[N];
	int fail[N],del[N],a[N],n,cnt;
	inline void update(int i){
		cnt++;
		int Mu=mul(dec(del[i],a[i]),ksm(dec(del[fail[cnt-2]],a[fail[cnt-2]]),mod-2));
		r[cnt].resize(i-fail[cnt-2],0);
		r[cnt].pb(Mu);
		for(int j=1;j<r[cnt-2].size();j++)
		r[cnt].pb(mul(Mu,mod-r[cnt-2][j]));
		r[cnt]=r[cnt]+r[cnt-1];
	}
	inline void BM(){
		for(int i=1;i<=n;i++){
			for(int j=1;j<r[cnt].size();j++)
			Add(del[i],mul(r[cnt][j],a[i-j]));
			if(del[i]!=a[i]){
				fail[cnt]=i;
				if(!cnt)r[++cnt].resize(i);
				else update(i);
			}
		}
		Cas::init(r[cnt]);
	}
	inline void init(int *v,int len){
		n=len;
		for(int i=1;i<=n;i++)a[i]=v[i-1];
		BM();
	}
}
vector<int> e[3005];
int f[3005][6005];
int n,K,rt,lim;
int main(){
	n=read(),lim=2*n;
	for(int i=1;i<n;i++){
		int u=read(),v=read();
		e[u].pb(v),e[v].pb(u);
	}
 	rt=read(),K=read();	
 	f[rt][0]=1;
 	for(int i=0;i<lim;i++)
 	for(int u=1;u<=n;u++)if(f[u][i])
 	for(int &v:e[u])Add(f[v][i+1],f[u][i]);
 	B_M::init(f[rt],lim);
 	poly res=Cas::calc(K);
 	for(int i=1;i<=n;i++){
 		int anc=0;
 		for(int j=0;j<res.size();j++)Add(anc,mul(f[i][j],res[j]));
 		cout<<anc<<" ";
	 }
}