AtCoder ABC #248 (ユニークビジョンプログラミングコンテスト2022) : G - GCD cost on the tree

Gまでは順調だったのに…。
https://atcoder.jp/contests/abc248/tasks/abc248_g

問題

木を成すグラフが与えられる。
各頂点vには、正整数値A[v]が設定されている。
2頂点間のコストを、両頂点を結ぶパス上の頂点に設定された整数値のGCDと、パス上の頂点数の積とする。
全頂点対におけるコストの総和を求めよ。

解法

f(d) := GCDがdとなるようなパスにおける頂点数の総和
とするとd*f(d)の総和が解となる。
f(d)を直接求める代わりに
g(d) := GCDがdの倍数となるようなパスにおける頂点数の総和
を求めれば、
f(d) = g(d) - sum(f(kd)) (kは2以上の整数)
より、f(d)を大きな順で求めることができる。

あとはdを総当たりしながらg(d)を求めよう。
dの倍数の整数値を持つ頂点と、そのような頂点を結ぶ辺それぞれからなる部分グラフについて、パス上の頂点数の総和をもとめればよい。
これは全方位木DPなどDFS2回で求めることができる。

int N;
int A[101010];
int U[101010],V[101010];
vector<int> E[101010];
vector<int> D[101010];
vector<int> cand[101010];

const ll mo=998244353;
ll num[101010];
int vis[202020];

int dfs(int cur,int id) {
	if(vis[cur]==id) return 0;
	vis[cur]=id;
	int v=1;
	FORR(e,E[cur]) v+=dfs(e,id);
	return v;
}

int dfs2(int cur,int pre,int n) {
	int v=0;
	FORR(e,E[cur]) if(e!=pre) {
		int a=dfs2(e,cur,n);
		v+=a;
		(num[vis[cur]]+=1LL*a*(n-v))%=mo;
	}
	v++;
	(num[vis[cur]]+=1LL*(n-v))%=mo;
	return v;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	for(i=1;i<=100000;i++) {
		for(j=i;j<=100000;j+=i) D[j].push_back(i);
	}
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	FOR(i,N-1) {
		cin>>x>>y;
		x--,y--;
		U[i]=x;
		V[i]=y;
		int g=__gcd(A[x],A[y]);
		FORR(v,D[g]) cand[v].push_back(i);
	}
	
	for(i=1;i<=100000;i++) {
		vector<int> C;
		FORR(e,cand[i]) {
			E[U[e]].push_back(V[e]);
			E[V[e]].push_back(U[e]);
			C.push_back(U[e]);
			C.push_back(V[e]);
		}
		FORR(c,C) if(vis[c]!=i) {
			x=dfs(c,i);
			dfs2(c,c,x);
		}
		FORR(e,cand[i]) {
			E[U[e]].clear();
			E[V[e]].clear();
		}
	}
	ll ret=0;
	for(i=100000;i>=1;i--) {
		for(j=2*i;j<=100000;j+=i) num[i]-=num[j];
		num[i]=(num[i]%mo+mo)%mo;
		ret+=num[i]*i%mo;
	}
	cout<<ret%mo<<endl;
	
}

まとめ

ちょっと実装が面倒だけど、本番ではまあまあの速さで解けたので良かった。