AT_abc376_g [ABC376G] Treasure Hunting 题解

考虑维护若干个联通快，初始化每个点一个连通块。对每个联通块维护大小、\(\sum_xa_x\) 还有最优期望，每次选择两个联通快合并。考虑两个联通快 \(a,b\) 哪个在前面更优。

记 \(a_c,a_s,a_d\) 分别表示连通块 \(a\) 的点数、\(\sum_xa_x\) 以及最优期望，\(b\) 同理。则 \(a\) 在 \(b\) 前更优当且仅当 \(a_c\cdot b_s<b_c\cdot a_s\)，用一个 set 维护即可，每次取出最优的连通块跟父亲合并，时间复杂度 \(\mathcal O(n\log n)\)。

参考代码：

#include<bits/stdc++.h>
#define ll long long
#define mxn 200003
#define md 998244353
#define rep(i,a,b) for(int i=(a);i<=(b);++i)
#define rept(i,a,b) for(int i=(a);i<(b);++i)
using namespace std;
ll power(ll x,int y){
	ll ans=1;
	for(;y;y>>=1){
		if(y&1)ans=ans*x%md;
		x=x*x%md;
	}
	return ans;
}
struct node{
	ll sz,ct,d;int x;
	inline bool operator<(node a)const{
		ll s1=ct*a.sz,s2=a.ct*sz;
		if(s1!=s2)return s1<s2;
		return x<a.x;
	}// c1*s2<c2*s1
}f[mxn];
inline node operator+(node x,node y){
	return {x.sz+y.sz,x.ct+y.ct,(x.d*x.sz+y.d*y.sz+x.ct*y.sz)%md*power(x.sz+y.sz,md-2)%md,x.x};
}
int T,n,fa[mxn],a[mxn],ft[mxn];
set<node>q;
inline int get(int x){
	return ft[x]==x?x:ft[x]=get(ft[x]);
}
void merge(int x,int y){
	x=get(x),y=get(y);
	q.erase(f[x]);
	ft[y]=x,f[x]=f[x]+f[y];
	q.insert(f[x]);
}
signed main(){
	scanf("%d",&T);
	while(T--){
		scanf("%d",&n);
		rep(i,0,n)ft[i]=i;
		rep(i,1,n)scanf("%d",&fa[i]);
		rep(i,1,n)scanf("%d",&a[i]);
		rep(i,0,n)f[i]={a[i],i>0,i>0,i},q.insert(f[i]);
		while(q.size()){
			int x=q.begin()->x;q.erase(q.begin());
			if(x)merge(fa[x],x);
		}
		cout<<f[get(0)].d<<'\n';
	}
	return 0;
}