【拉格朗日优化dp】P4365 [九省联考 2018] 秘密袭击 coat

题目简述

求树上所有连通块第 $k$ 大点权（不足 $k$ 点记为 $0$）的和。

$1\leq k\leq n\leq1666$，$d_i\leq W\leq 1666$。

解题思路

我们对第 $k$ 大的和的这种问题，通常的处理方法是对值域枚举一个阈值 $i$。对于 $\geq i$ 的点视作为 $1$，$< i$ 的点视作为 $0$。这样子问题就变成了了所有 $i$，连通块第 $k$ 大是 $1$（也就是有 $\geq k$ 个 $1$）的数目的和。

假设 $f[u][j]$ 表示连通块以 $u$ 为根且有 $j$ 个 $1$ 节点，我们可以把问题转成树上背包。

通过经典结论，合理控制枚举范围我们可以做到 $O(n^2)$ 来转移，总得时间复杂度 $O(Wn^2)$，随便常数优化一下跑得比正解还要快。

我们下面讨论正解。

众所周知的结论，树上背包可以转为生成函数，也就是说假设 $F_u(x)=f[u][0]+f[u][1]x+f[u][2]x^2+\dots$。那么就会有：

\[F_u(x)=x^{d[u]\geq i}\prod_{v\in u} (F_v(x)+1)
\]

多项式卷积很慢，而且不便于转移，所以我们不应该卷积，考虑点值表示法。

假设我们选取的点为 $t$，那么转移变为：

\[F_u=t^{d[u]\geq i}\prod_{v\in u}(F_v+1)
\]

不难发现阈值 $i$ 对整个 $F$ 的影响比较小，我们不妨考虑换一种思考方式，也就是枚举点 $t$，然后快速搞定每一个阈值。我们期望每一次枚举 $t$ 可以做到 $O(n\log W)$ 得到所有阈值的值。

我们考虑每个位置维护一个当前值数组 $val$，那么每一个节点相当于是他所有叶节点的对应 $val$ 位 $+1$ 再相乘，然后再给 $[1,d[u]]$ 的位乘上 $t$。暴力乘积很慢，但是可以发现我们需要维护的东西是：

全局加。
区间乘。
合并两个 $val$，对应位相乘。

这个结构是很显然的线段树合并，这里就可以做到 $O(n\log W)$ 得到所有阈值了。

还没完，有一个小的问题，就是一个连通块的根不是 $1$ 时由于线段树合并了，所以他的 $val$ 会被他的祖先修改掉导致不可用，解决这个问题，我们需要一个辅助生成函数 $G_u(x)$ 表示所有 $u$ 的子树节点的 $val$，那我们的 $G(x)$ 需要维护一个：

合并两个 $G$，对应位相加。
把 $F$ 的对应位相加到 $G$。

感觉对 $G$ 单独再开个动态开点线段树可以做，但是写起来超级麻烦，介绍一个相对简单的方法。

考虑矩阵乘法：

\[\begin{bmatrix}f&g&1\end{bmatrix}\times \begin{bmatrix}a&b&0\\0 &1&0\\c&d&1\end{bmatrix}=\begin{bmatrix}af+c&bf+g+d&1\end{bmatrix}
\]

矩阵 $\begin{bmatrix}a&b&0\\0 &1&0\\c&d&1\end{bmatrix}$ 具有可乘形式，具体来说：

\[\begin{bmatrix}a&b&0\\0 &1&0\\c&d&1\end{bmatrix}\times \begin{bmatrix}a'&b'&0\\0 &1&0\\c'&d'&1\end{bmatrix}=\begin{bmatrix}aa'&ab'+b&0\\0 &1&0\\ca'+c'&cb'+d+d'&1\end{bmatrix}
\]

我们可以维护这种矩阵，可以发现这种矩阵可以维护：

全局加 $1$：$a=c=1,b=d=0$。
把 $f$ 加到 $g$ 上：$a=b=1,c=d=0$。

其他操作需要线段树合并。我们就完成了本题的求点值部分。

对于转回多项式，我们有拉格朗日插值法：

\[f(x)=\sum_{i=1}^n y_i\prod_{j\not =i} \dfrac{x-j}{i-j}
\]

不难发现这个式子每个 $i$ 的常数可以先预处理出来，多项式卷积部分我们可以先求暴力 $\prod(x-j)$ 然后每次除 $x-i$ 即可，暴力除是 $O(n)$ 的，原不及瓶颈线段树合并的大常数 $O(n^2\log n)$。

参考代码

#include<iostream>

#include<vector>

using namespace std;

#define ll long long

const int MOD=64123;

const int MAXN=1667;

int n,k,w;

ll d[MAXN],f[MAXN];

ll fac[MAXN],inf[MAXN];

int rt[MAXN];

vector<int> ve[MAXN];

int add(int x){return (x%MOD+MOD)%MOD;}

ll inv(ll a,int b=MOD-2){ll res=1;while(b){if(b&1) res=res*a%MOD;a=a*a%MOD;b>>=1;}return res;}

struct matrix{

	ll a,b,c,d;

	matrix(ll aa=0,ll bb=0,ll cc=0,ll dd=0){a=aa,b=bb,c=cc,d=dd;}

	matrix operator *(matrix i){return matrix(1ll*a*i.a%MOD,add(1ll*i.a*b%MOD+i.b),

		add(1ll*a*i.c%MOD+c),add(1ll*b*i.c%MOD+add(d+i.d)));}

	bool operator ==(matrix i){return (i.a==a)&&(i.b==b)&&(i.c==c)&&(i.d==d);}

};

matrix I=(matrix){1,0,0,0};

struct node{

	int ls,rs;

	matrix val;

}t[MAXN*40];int tot=0;

struct poly{

vector<ll> ve;

void init(int s){ve.resize(s);}

poly operator +(poly x){

	poly y;y.init(max(x.ve.size(),ve.size()));

	for(int i=0;i<y.ve.size();i++)

		y.ve[i]=add(((i<ve.size())?ve[i]:0)+

				((i<x.ve.size())?x.ve[i]:0));

	return y;

}

poly operator *(ll x){

	poly y;y.init(ve.size());

	for(int i=0;i<ve.size();i++)

		y.ve[i]=ve[i]*x%MOD;

	return y;

}

poly operator *(poly x){

	poly y;y.init(x.ve.size()+ve.size()-1);

	for(int i=0;i<x.ve.size();i++)

		for(int j=0;j<ve.size();j++)

			y.ve[i+j]=add(y.ve[i+j]+x.ve[i]*ve[j]%MOD);

	return y;

}

poly operator /(ll x){

	poly y;y.init(ve.size()-1);

	ll res=ve[ve.size()-1];

	for(int i=ve.size()-2;i>=0;i--){

		y.ve[i]=res;

		res=add(ve[i]+x*res%MOD);

	}

	return y;

}

};

void newnode(int &k){k=++tot;t[k]=node{0,0,I};}

void pushdown(int k){

	if(t[k].val==I) return;

	if(!t[k].ls) newnode(t[k].ls);

	if(!t[k].rs) newnode(t[k].rs);

	t[t[k].ls].val=t[t[k].ls].val*t[k].val;

	t[t[k].rs].val=t[t[k].rs].val*t[k].val;

	t[k].val=I;

}

void modify(int &id,int l,int r,int L,int R,matrix val){

	if(l>R||r<L) return;

	if(!id) newnode(id);

	if(L<=l&&r<=R) {t[id].val=t[id].val*val;return;}

	int mid=l+r>>1;

	pushdown(id);

	modify(t[id].ls,l,mid,L,R,val);

	modify(t[id].rs,mid+1,r,L,R,val);

}

void merge(int &x,int y,int l,int r){

	if(!t[x].ls&&!t[x].rs) swap(x,y);

	if(!t[y].ls&&!t[y].rs){

		t[x].val=t[x].val*matrix{t[y].val.b,0,0,t[y].val.d};

		return;

	}

	int mid=l+r>>1;

	pushdown(x);pushdown(y);

	merge(t[x].ls,t[y].ls,l,mid);

	merge(t[x].rs,t[y].rs,mid+1,r);

}

void dfs(int u,int fa,int x){

	modify(rt[u],1,w,1,d[u],matrix{0,x,0,0});

	modify(rt[u],1,w,d[u]+1,w,matrix{0,1,0,0});

	for(int v:ve[u]) if(v!=fa) {

		dfs(v,u,x);

		modify(rt[v],1,w,1,w,matrix{1,1,0,0});

		merge(rt[u],rt[v],1,w);

	}

	modify(rt[u],1,w,1,w,matrix{1,0,1,0});

	return;

}

int query(int id,int l,int r){

	if(l==r) return t[id].val.d;

	int mid=l+r>>1;

	pushdown(id);

	return add(query(t[id].ls,l,mid)+query(t[id].rs,mid+1,r));

}

signed main(){

	cin>>n>>k>>w;

	for(int i=1;i<=n;i++)

		cin>>d[i];

	for(int i=1;i<n;i++){

		int u,v;cin>>u>>v;

		ve[u].push_back(v);

		ve[v].push_back(u);

	}

	for(int i=0;i<=n;i++){

		tot=0;

		for(int j=1;j<=n;j++) rt[j]=0;

		dfs(1,0,i);

		f[i]=query(rt[1],1,w);

	}

	poly p;p.init(1);p.ve[0]=1;

	poly v;v.init(2);v.ve[1]=1;

	for(int i=0;i<=n;i++){

		v.ve[0]=(MOD-i)%MOD;

		p=p*v;

	}

	poly F;F.init(n+1);

	fac[0]=inf[0]=1;

	for(int i=1;i<=n;i++)

		fac[i]=fac[i-1]*i%MOD,inf[i]=inv(fac[i]);

	for(int i=0;i<=n;i++){

		int val=1ll*f[i]*inf[i]%MOD*inf[n-i]%MOD;

		if((n-i)&1) val=MOD-val;

		poly y=(p/i)*val;

		F=F+y;

	}

	int res=0;

	for(int i=k;i<=n;i++)

		res=add(res+F.ve[i]);

	cout<<res;

	return 0;

}