洛谷 P6834 [Cnoi2020] 梦原 题解

Announcement

• Programmed on 2024/2/25
• Written on 2024/2/26

题目来源

• 洛谷 P6834

Description

有一棵有 $n$ 个节点、根节点编号为 $1$、形态不确定的树，且每个点都有个权值 $a_i$。

对于任意节点 $i\in [2,n]\cap \mathbb{Z}$，其父亲将从 $[i-k,i-1]\cap {\mathbb{N}}^{+}$ 中等概率选取一个。

当树的形态确定后，将进行若干次操作，每次操作可以选择节点权值均不为 $0$ 的一个联通块，并将该联通块内每个节点的权值减去 $1$。

求最少操作次数的期望值（$\mathrm{mod}998244353$）。

• $1\le k<n\le 10^6$。

Solution

考虑树的形态确定时，最少的操作次数 $ans$。

类似于 洛谷 P1969 [NOIP2013 提高组] 积木大赛 的树上做法：
$$ans=\sum _{i=1}^n\frac{{\sum }_{j=1}^{\min (i-1,k)}\max (a_i-a_{i-j},0)}{\min (i-1,k)}$$
若暴力统计效率为 $O(nk)$，需要优化。

观察发现，计算 $\sum _{j=1}^{\min (i-1,k)}\max (a_i-a_{i-j},0)$ 时存在复杂度瓶颈。

由于只有当 $a_i>a_{i-j}$ 时才对答案有贡献，则该算式等价于 $\sum _{j=1}^{\min (i-1,k)}[a_i>a_{i-j}]a_i-[a_i>a_{i-j}]a_{i-j}$。

那么，我们只需要统计所有权值小于 $a_i$ 的节点个数，以及对应的节点权值和。

于是可以开两个树状数组分别统计满足条件的节点个数和权值和。

注意到若一个节点 $x<i-k$ 则需从树状数组中移除。

再加上离散化，复杂度可以达到 $O(n\log n)$。

Conclusion

• 对于统计 $i$ 的某些权值时只需统计满足 $a_j<a_i$ 的 $j$ 的 $a_j$ 之和等类似问题，可以考虑树状数组优化。
• 若只关心各权值的大小关系而不关心具体数值，考虑离散化技巧，进而避免 MLE / TLE 等。

Code

#include <bits/stdc++.h>
using namespace std;
const int p=998244353;
int n,k,a[1000005],b[1000005],f[1000005],c[2][1000005],tot; // c[0]-权值 c[1]-个数
unordered_map<int,int> id;
int qpow(int a,int b){
	int res=1;
	for (;b;b>>=1,a=1ll*a*a%p) if (b&1) res=1ll*res*a%p;
	return res;
}
void update(int op,int x,int d){
	for (;x<=tot;x+=x&(-x)) (c[op][x]+=d)%=p;
}
int query(int op,int x){
	int res=0;
	for (;x;x-=x&(-x)) (res+=c[op][x])%=p;
	return res;
}
int main(){
	scanf("%d%d",&n,&k);
	for (int i=1;i<=n;i++) scanf("%d",&a[i]),b[i]=a[i];
	sort(b+1,b+n+1);
	for (int i=1;i<=n;i++) if (i==n||b[i]!=b[i+1]) id[b[i]]=++tot;
	f[1]=a[1],update(0,id[a[1]],a[1]),update(1,id[a[1]],1);
	for (int i=2;i<=n;i++){
		update(0,id[a[i]],a[i]),update(1,id[a[i]],1);
		f[i]=(f[i-1]+1ll*qpow(min(i-1,k),p-2)*((1ll*query(1,id[a[i]]-1)*a[i]-query(0,id[a[i]]-1))%p+p)%p)%p;
		if (i-k>0) update(0,id[a[i-k]],-a[i-k]),update(1,id[a[i-k]],-1);
	}
	printf("%d\n",f[n]);
	return 0;
}