CSAcademy Round #58 : E. Path Inversions

最初の発想がなかったのでどうしようもない。
https://csacademy.com/contest/round-58/task/path-inversions/

問題

木を成す無向グラフが与えられる。
この木において、同じ点・辺を複数通らない長さKのパスを考えよう。
パス上の頂点番号を並べたときのinversionの総和を求めよ。

解法

inversionを考えるのが面倒なため、まずここを考えよう。
始点と終点と逆にしただけの２つのパスを考えると、パス上の２頂点は２つのパスのいずれか片方で必ずinversionが発生する。
よって、向きを無視してパスを考えるとパス1本あたり必ず個のinversionが発生する。

あとは向きを無視して長さKのパスの数を数えればよい。
これは重心分解を行いつつ、重心からの距離のカウントを保持しつつ数え上げていけばよい。

int N,K;
vector<int> EV[101010];
ll ret;
ll mo=1000000007;
int vis[101010];
int NV[101010];

int cnt[101010];
vector<int> C,D;

int dfs(int cur,int pre) {
	NV[cur]=1;
	FORR(e,EV[cur]) if(e!=pre && vis[e]==0) NV[cur]+=dfs(e,cur);
	return NV[cur];
}

int dfs2(int cur,int pre,int TNV) {
	int tot=1;
	int ok=1;
	FORR(e,EV[cur]) if(e!=pre && vis[e]==0) {
		int c = dfs2(e,cur,TNV);
		if(c!=-1) return c;
		tot += NV[e];
		if(NV[e]*2>TNV) ok=0;
	}
	if((TNV-tot)*2>TNV) ok=0;
	if(ok) return cur;
	return -1;
}

void dfs3(int cur,int pre,int d) {
	if(d<=K) {
		(ret+=cnt[K-d])%=mo;
	}
	FORR(e,EV[cur]) if(vis[e]==0 && e!=pre) dfs3(e,cur,d+1);
	C.push_back(d);
}

void split(int cur,int id) {
	int TNV = dfs(cur,-1);
	if(TNV==0) return;
	int center=dfs2(cur,-1,TNV);
	cnt[0]++;
	FORR(e,EV[center]) if(vis[e]==0) {
		C.clear();
		dfs3(e,center,1);
		FORR(v,C) cnt[v]++, D.push_back(v);
	}
	cnt[0]--;
	FORR(v,D) cnt[v]--;
	D.clear();
	vis[center]=1;
	FORR(e,EV[center]) if(vis[e]==0) split(e,id+1);
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N-1) {
		cin>>x>>y;
		EV[x-1].push_back(y-1);
		EV[y-1].push_back(x-1);
	}
	split(0,0);
	cout<<(ret*(((1LL*K*(K+1))/2)%mo))%mo<<endl;
	
}

まとめ

inversionを最初に片づける発想がなぜ思いつかなかった…。