AtCoder ABC #359 (ユニークビジョンプログラミングコンテスト2024 夏) : G - Sum of Tree Distance

想定解と異なるためか、TLE連発してしまった。
https://atcoder.jp/contests/abc359/tasks/abc359_g

問題

N頂点の木を成す無向グラフが与えられる。
各点には色が設定されている。
同色の頂点対の最短距離の総和を求めよ。

解法

自分は平方分割で解いた。
あらかじめLCAを算出できるようにしておき、2点間の距離をO(logN)で求められるようにしておこう。

以下、色ごとに処理する。

• 同色の頂点数が少ない場合
  • 頂点対を総当たりし、LCAを用いてその距離を求める。
• 同色の頂点数が多い場合
  • 各辺が解に計上されるのは、辺の両側のその色の頂点数の積である。
  • よって、DFSしながら辺の両端のその色の頂点数を求めて行こう。
  • ただし、自分の場合DFSするとTLEしたので別の方法で、辺の両端の頂点数を求めた。DFS順の頂点訪問順と、それにおける区間内の同色の頂点数の累積和を取ればよい。

int N;
vector<int> E[200005];
int P[21][200005],D[200005];
int id;
ll ret=0;
int A[202020];

int L[202020],R[202020],re[202020];
int S[202020];

void dfs(int cur) {
	L[cur]=id++;
	re[L[cur]]=cur;
	FORR(e,E[cur]) if(e!=P[0][cur]) D[e]=D[cur]+1, P[0][e]=cur, dfs(e);
	R[cur]=id;
}
vector<int> V[202020];

int dist(int a,int b) {
	int ret=0,i,aa=a,bb=b;
	if(D[aa]>D[bb]) swap(aa,bb);
	for(i=17;i>=0;i--) if(D[bb]-D[aa]>=1<<i) bb=P[i][bb];
	for(i=17;i>=0;i--) if(P[i][aa]!=P[i][bb]) aa=P[i][aa], bb=P[i][bb];
	return D[a]+D[b]-2*D[(aa==bb)?aa:P[0][aa]];  // dist
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0);
	FOR(i,17) FOR(x,N) P[i+1][x]=P[i][P[i][x]];
	FOR(i,N) {
		cin>>A[i];
		A[i]--;
		V[A[i]].push_back(i);
	}
	
	FOR(i,N) {
		if(V[i].size()<=100) {
			FORR(x,V[i]) FORR(y,V[i]) if(x<y) ret+=dist(x,y);
		}
		else {
			FOR(j,N) {
				S[j+1]=S[j]+(A[re[j]]==i);
			}
			FOR(j,N) {
				x=re[j];
				y=S[R[j]]-S[L[j]];
				ret+=1LL*(V[i].size()-y)*y;
			}
		}
	}
	cout<<ret<<endl;
}

まとめ

うーむ、Writer想定解ではないのね。