LOJ#6289. 花朵 树链剖分+分治NTT
本来以为这道题会非常难调，但是没想到调了不到 5 分钟就 A 了.
由于基于多项式的运算都可以方便地进行封装，所以细节就不是很多（或者说几乎没有细节）
题意：给定一棵树，每个点有点权，求对于所有大小为 $m$ 的独立集的点权之积的和.
数据范围：$n,m \leqslant 8 \times 10^4$.
先考虑一个十分显然的 $O(n^2)$ 暴力：
令 $f[x][i],g[x][i]$ 分别表示点 $x$ 选/不选的情况下独立集大小为 $i$ 的点积 之和.
考虑将 $x$ 与 $x$ 的一个儿子 $y$ 合并：$f[x][i+j]=f[x][i] \times f[y][j]$，$g$ 同理.
然后 $x$ 的初始值是：$f[x][1]=w[x],g[x][0]=1$.
树形DP 卡一下上界复杂度是 $O(n^2)$ 的.
不难发现，上述 $f[x][i+j] = f[x][i] \times f[y][j]$ 是一个卷积的形式.
如果是菊花图或者链的话可以直接用 NTT/分治NTT 来做.
正解的话考虑进行轻重路径剖分：
对于一条重链来说，先求出该重链中每个点轻儿子为根的多项式 $f,g$，然后对于重链中每个点都将其轻儿子与该点合并.
最后对于一条重链进行分治，求出该重链链顶为根的多项式.
分析一下时间复杂度：
考虑一条重链链顶为根的子树会被卷多少次：其祖先中每一条重链都会将其贡献一次.
那么树链剖分中一个点有 $O(\log n)$ 个祖先，而每次卷积的时候对链分治的复杂度是 $O(n \log^2 n)$.
总复杂度就是 $O(n \log^3 n)$，但是由于树链剖分的常数比较小，跑的并不慢.
code:
#include <queue>
#include <cstdio>
#include <vector>
#include <cstring>
#include <algorithm>
#define N 1000009
#define ll long long
#define mod 998244353
#define pb push_back
#define setIO(s) freopen(s".in","r",stdin)
using namespace std;
int m;
int A[N<<2],B[N<<2];
int tim,edges,n;
int size[N],son[N],top[N],hd[N],to[N<<1],nex[N<<1],fa[N],dep[N];
int dfn[N],bu[N],si[N],val[N];
void add(int u,int v) {
nex[++edges]=hd[u],hd[u]=edges,to[edges]=v;
}
int ADD(int x,int y) {
return (ll)(x+y)%mod;
}
int DEC(int x,int y) {
return (ll)(x-y+mod)%mod;
}
int MUL(int x,int y) {
return (ll)x*y%mod;
}
int qpow(int x,int y) {
int tmp=1;
for(;y;y>>=1,x=(ll)x*x%mod) {
if(y&1) tmp=(ll)tmp*x%mod;
}
return tmp;
}
int get_inv(int x) {
return qpow(x,mod-2);
}
void NTT(int *a,int len,int op) {
for(int i=0,k=0;i<len;++i) {
if(i>k) {
swap(a[i],a[k]);
}
for(int j=len>>1;(k^=j)<j;j>>=1);
}
for(int l=1;l<len;l<<=1) {
int wn=qpow(3,(mod-1)/(l<<1));
if(op==-1) wn=get_inv(wn);
for(int i=0;i<len;i+=l<<1) {
int w=1;
for(int j=0;j<l;++j) {
int x=a[i+j],y=(ll)w*a[i+j+l]%mod;
a[i+j]=(ll)(x+y)%mod;
a[i+j+l]=(ll)(x-y+mod)%mod;
w=(ll)w*wn%mod;
}
}
}
if(op==-1) {
int iv=get_inv(len);
for(int i=0;i<len;++i) {
a[i]=(ll)a[i]*iv%mod;
}
}
}
struct poly {
int len;
vector<int>a;
poly() { len=0,a.clear(); }
void push(int x) {
a.pb(x),++len;
}
void resize(int x) {
a.resize(x),len=x;
}
poly operator*(const poly &b) const {
int lim;
for(lim=1;lim<len+b.len-1;lim<<=1);
for(int i=0;i<lim;++i) A[i]=B[i]=0;
for(int i=0;i<len;++i) A[i]=a[i];
for(int i=0;i<b.len;++i) B[i]=b.a[i];
NTT(A,lim,1),NTT(B,lim,1);
for(int i=0;i<lim;++i) {
A[i]=(ll)A[i]*B[i]%mod;
}
NTT(A,lim,-1);
poly c;
for(int i=0;i<len+b.len-1;++i) {
c.push(A[i]);
}
if(c.len>m+1) c.resize(m+1);
return c;
}
poly operator+(const poly &b) const {
poly c;
c.resize(max(len,b.len));
for(int i=0;i<c.len;++i) c.a[i]=0;
for(int i=0;i<c.len;++i) {
if(i<len) c.a[i]=ADD(c.a[i],a[i]);
if(i<b.len) c.a[i]=ADD(c.a[i],b.a[i]);
}
return c;
}
poly operator-(const poly &b) const {
poly c;
c.resize(max(len,b.len));
for(int i=0;i<c.len;++i) c.a[i]=0;
for(int i=0;i<c.len;++i) {
if(i<len) c.a[i]=ADD(c.a[i],a[i]);
if(i<b.len) c.a[i]=DEC(c.a[i],b.a[i]);
}
return c;
}
}f0[N],f1[N],g[2][N];
struct data {
poly f00,f01,f10,f11;
data operator+(const data &b) const {
data c;
c.f00=(f01*b.f00)+(f00*(b.f00+b.f10));
c.f11=(f11*b.f01)+(f10*(b.f11+b.f01));
c.f01=(f01*b.f01)+(f00*(b.f01+b.f11));
c.f10=(f11*b.f00)+(f10*(b.f10+b.f00));
return c;
}
}tmp;
void dfs1(int x,int ff) {
fa[x]=ff,dep[x]=dep[ff]+1,size[x]=1;
for(int i=hd[x];i;i=nex[i]) {
int y=to[i];
if(y==ff) continue;
dfs1(y,x);
size[x]+=size[y];
if(size[y]>size[son[x]]) son[x]=y;
}
}
void dfs2(int x,int tp) {
top[x]=tp;
dfn[x]=++tim;
bu[tim]=x;
++si[tp];
if(son[x]) {
dfs2(son[x],tp);
}
for(int i=hd[x];i;i=nex[i]) {
if(to[i]!=fa[x]&&to[i]!=son[x]) {
dfs2(to[i],to[i]);
}
}
}
poly calc(int l,int r,int d) {
if(l==r) {
return g[d][l];
}
int mid=(l+r)>>1;
return calc(l,mid,d)*calc(mid+1,r,d);
}
data solve(int l,int r) {
if(l==r) {
int u=bu[l];
data e;
e.f00=f0[u];
e.f11=f1[u];
return e;
}
int mid=(l+r)>>1;
return solve(l,mid)+solve(mid+1,r);
}
int main() {
// setIO("input");
int x,y,z;
scanf("%d%d",&n,&m);
for(int i=1;i<=n;++i) scanf("%d",&val[i]);
for(int i=1;i<n;++i) {
scanf("%d%d",&x,&y);
add(x,y),add(y,x);
}
dfs1(1,0),dfs2(1,1);
for(int i=1;i<=n;++i) {
f0[i].push(1);
f1[i].push(0);
f1[i].push(val[i]);
}
for(int i=n;i>=1;--i) {
int p=bu[i];
if(top[p]==p) {
for(int j=dfn[p];j<=dfn[p]+si[p]-1;++j) {
x=bu[j];
int p0=0,p1=0;
for(int e=hd[x];e;e=nex[e]) {
y=to[e];
if(y==son[x]||y==fa[x]) continue;
g[0][++p0]=f0[y]+f1[y];
g[1][++p1]=f0[y];
}
if(p0) f0[x]=calc(1,p0,0);
if(p1) f1[x]=f1[x]*calc(1,p1,1);
}
tmp=solve(dfn[p],dfn[p]+si[p]-1);
f0[p]=tmp.f01+tmp.f00;
f1[p]=tmp.f10+tmp.f11;
}
}
f0[1].resize(m+1);
f1[1].resize(m+1);
printf("%d\n",(ll)(f0[1].a[m]+f1[1].a[m])%mod);
return 0;
}
