AtCoder ABC #376 (Promotion of AtCoder Career Design DAY) : G - Treasure Hunting

最近このテク多いのか…。
https://atcoder.jp/contests/abc376/tasks/abc376_g

問題

根付き木が与えられる。
いずれかの点にお宝があり、それぞれの点にお宝がある確率が与えられる。

初期状態で根頂点だけ訪問している状態である。
未訪問の点のうち、親頂点を訪問済みの点をどこか選択して訪問するとする。
訪問回数が最少となるように最適な手順で訪問したとき、お宝のある点に至るまでの訪問頂点数の期待値を求めよ。

解法

頂点集合vに対し、A(v)を集合内の確率の総和、C(v)を集合内の頂点数とする。
A(v)/C(v)が最大のものは、親頂点を含む頂点集合の直後に訪問する、としてよい。
この処理は、頂点集合内の最も根に近い頂点と最後に訪問した頂点を持って置き、union_findで連結していけばよい。
その際、親頂点の集合をpとするとA(p)*C(v)の総和を解に計上する。

int T;
int N;
int P[202020];
const ll mo=998244353;
int A[202020];
int C[202020];
ll S[202020];

int tail[202020];
int nex[202020];

template<int um> class UF {
	public:
	vector<int> par,rank,cnt,G[um];
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit(int num=um) {int i; FOR(i,num) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(x<y) return par[y]=x;
		else return par[x]=y;
	}
};
UF<202020> uf;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}


struct Edge {
	ll S,C;
	int id;
	bool operator<(const Edge &y) const {
		return (S*y.C > y.S*C)||((S*y.C == y.S*C&&id<y.id));
	}
};

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		N++;
		for(i=1;i<N;i++) {
			cin>>P[i];
		}
		ll sum=0;
		uf.reinit(N);
		set<Edge> Q;
		
		for(i=0;i<N;i++) {
			if(i) cin>>A[i];
			sum+=A[i];
			S[i]=A[i];
			C[i]=1;
			tail[i]=i;
			if(i) Q.insert({S[i],C[i],i});
		}
		
		ll ret=0;
		while(Q.size()) {
			x=Q.begin()->id;
			Q.erase(Q.begin());
			
			y=uf[P[x]];
			(ret+=1LL*C[y]*S[x])%=mo;
			Q.erase({S[y],C[y],y});
			S[y]+=S[x];
			C[y]+=C[x];
			if(y) {
				Q.insert({S[y],C[y],y});
			}
			
			uf(x,y);
		}
		
		ret=ret%mo*modpow(sum)%mo;
		cout<<ret<<endl;
		
	}
}

まとめ

これAGCのボス問級の問題だったのか…。