Codeforces #537 Div2 E. Tree

考察は難しくないが、実装が結構手間。
http://codeforces.com/contest/1111/problem/E

問題

木を成す無向グラフが与えられる。
このグラフに対し、以下のクエリに答えよ。

• 根頂点Rの他、K個の頂点が指定される。
• このK個の頂点群を、M個以下のグループに分ける。ただし互いに祖先・子孫の関係にある頂点対は、同じグループに入れてはならない。
• この時、グループの分け方は何通りか。

解法

頂点Rに近い順に、K個の頂点を順次処理しよう。
dp(n,m) := n個目までの頂点のグループ分けを決めたとき、m個に分かれる場合の組み合わせ
とする。dp(n,m)がわかっているとき、(n+1)個目の頂点vの分け方を考える。
K個の頂点群の中に、vの祖先にあたる頂点がp個あるとする。
その場合、m個中p個のグループには属せない。よって残り(m-p)個のどこかに属するか、新規グループを作るかの選択肢が残る。
よって、dp(n+1,m) += dp(n,m)*(m-p)、dp(n+1,m+1) += dp(n,m)でテーブルを更新しよう。

「頂点Rに近い順に、K個の頂点を並べる」のは、先に根を決め打ちした状態でLCAを求められるようにしておけば、各頂点とRとの距離が高速に求められるのでまぁ問題ない。
次に、「vの祖先にあたる頂点がp個あるとする。」のpを求めなければならない。
これはEulerTour+BITで、頂点vを1個処理するたびにBITでSubTree内に1加算することで対応しよう。
ただしこの問題では根頂点Rがコロコロ変わるので、Rが元のグラフでSubTreeにある場合、Rから見たvのSubTreeはEulerTour順で２つの区間になる場合がある点に注意。

int N,Q;
vector<int> E[200005];
int P[21][200005],D[200005];
int L[200005],R[200005],eid;
int K,M,RR;

ll from[303];
ll to[303];
ll mo=1000000007;

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<int,20> bt;

void dfs(int cur) {
	L[cur]=eid++;
	ITR(it,E[cur]) if(*it!=P[0][cur]) D[*it]=D[cur]+1, P[0][*it]=cur, dfs(*it);
	R[cur]=eid;
}
int getpar(int cur,int up) {
	int i;
	FOR(i,20) if(up&(1<<i)) cur=P[i][cur];
	return cur;
}

pair<int,int> lcadist(int a,int b) {
	int ret=0,i,aa=a,bb=b;
	if(D[aa]>D[bb]) swap(aa,bb);
	for(i=19;i>=0;i--) if(D[bb]-D[aa]>=1<<i) bb=P[i][bb];
	for(i=19;i>=0;i--) if(P[i][aa]!=P[i][bb]) aa=P[i][aa], bb=P[i][bb];
	return make_pair((aa==bb)?aa:P[0][aa], D[a]+D[b]-2*D[(aa==bb)?aa:P[0][aa]]);
}
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	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,19) FOR(x,N) P[i+1][x]=P[i][P[i][x]];
		
	while(Q--) {
		cin>>K>>M>>RR;
		RR--;
		vector<vector<int>> V;
		FOR(i,K) {
			cin>>r;
			r--;
			auto p=lcadist(RR,r);
			V.push_back({p.second,p.first,r});
		}
		
		ZERO(from);
		from[0]=1;
		sort(ALL(V));
		FORR(v,V) {
			int cur=v[2];
			int p=bt(L[cur]);
			ZERO(to);
			FOR(i,M+1) {
				if(i>=p) (to[i]+=from[i]*(i-p))%=mo;
				(to[i+1]+=from[i])%=mo;
			}
			if(cur==RR) {
				bt.add(0,1);
			}
			else if(cur==v[1]) {
				bt.add(0,1);
				p=getpar(RR,D[RR]-D[cur]-1);
				bt.add(L[p],-1);
				bt.add(R[p],1);
			}
			else {
				bt.add(L[cur],1);
				bt.add(R[cur],-1);
			}
			swap(from,to);
		}
		ll ret=0;
		FOR(i,M+1) ret+=from[i];
		cout<<(ret%mo)<<endl;
		FORR(v,V) {
			int cur=v[2];
			if(cur==RR) {
				bt.add(0,-1);
			}
			else if(cur==v[1]) {
				int p=getpar(RR,D[RR]-D[cur]-1);
				bt.add(0,-1);
				bt.add(L[p],+1);
				bt.add(R[p],-1);
			}
			else {
				bt.add(L[cur],-1);
				bt.add(R[cur],1);
			}
		}
	}
	
}

まとめ

割とややこしいけど、すんなり解けて良かった。
実装はともかく考察自体は軽いので、Div2の最終問題っぽい感じ。