HourRank20 : C. Birjik and Nicole's Tree Game

このテクは勉強になった。
https://www.hackerrank.com/contests/hourrank-20/challenges/birjik-and-nicoles-tree-game

問題

N頂点からなる根付き木が与えられる。
この木に対する以下のQ個のクエリに答えよ。

N頂点中、K個の頂点を選び黒く塗る。他の頂点は白とする。
以下のように定義されるC(i)について、C(0)～C(K)の値を求めよ。
C(i) := 頂点vのうち、vのSubTree中に黒頂点がちょうどi個あるような者の数

解法

N頂点すべてを考えるのではなく、黒頂点およびそれらのLCAだけで構成される頂点群および辺を考える。
SubTreeの黒頂点数が変化するのは、上記頂点群だけである。
よってその間の黒でもLCAでもない頂点群はまとめて数えてしまおう。

「黒頂点およびそれらのLCAだけで構成される頂点群」は下記の通り構成すると、元の黒頂点の数の倍以下にしかならない。
全クエリで処理すべき頂点数はO(sum(K))で済む。

• まず全頂点vをDFSで訪問し訪問順に番号を振る。(ついでに深さD(v)も求めておく)
• 黒頂点群を上記訪問順でソートし、隣接する黒頂点間のLCAをもとの黒頂点集合に加える

最後の操作はLCAの頂点が集合に加わったからと言って複数回行う必要はなく、元々ある黒頂点集合間でのみ行えばよい。

あとはこうして構成した木をDFSしつつ数え上げる。
辺を再構成してもよいのだが、先ほどの訪問順ソートを流用すると、明に辺をもたずにDFSできる。

根頂点から現在いる頂点までの経路上の頂点を並べたスタックSを考える。
訪問順ソート順に頂点vを処理したとするとき、以下の処理を行えばよい。

• スタックの末尾の頂点xが頂点vの祖先でない限り、xを取り除く
  • その際、末尾の頂点x～末尾から2番目の頂点yの間のD(x)-D(y)個の頂点は、xと同じ黒頂点をSubTreeに持つので、その分C(i)に加算する。
• スタックにvを追加する。

全頂点探索後、Sの要素が２個以上ある間、Sの末尾を取り除き、「その際、末尾の頂点x～～」の処理を行う。

int N;
vector<int> E[303030];
vector<pair<int,int>> MD;
int Q,K;

int P[21][300005],D[300005];
int ID[303030],id;
int num[303030];
int ret[303030];


void dfs(int cur,int pre,int d) {
	D[cur]=d;
	P[0][cur]=pre;
	ID[cur]=++id;
	
	int i,p=-1;
	FOR(i,E[cur].size()) {
		if(E[cur][i]==pre) p=i;
		else dfs(E[cur][i],cur,d+1);
	}
	if(p>=0) E[cur].erase(E[cur].begin()+p);
}
int lca(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 (aa==bb)?aa:P[0][aa];               // vertex
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x].push_back(y);
		E[y].push_back(x);
	}
	
	dfs(1,0,1);
	FOR(i,19) FOR(x,N+1) P[i+1][x]=P[i][P[i][x]];
	
	cin>>Q;
	while(Q--) {
		cin>>K;
		
		vector<pair<int,int>> V;
		V.push_back({0,0});
		FOR(i,K) {
			cin>>x;
			num[x]=1;
			V.push_back({ID[x],x});
		}
		sort(ALL(V));
		FOR(i,K) {
			x = lca(V[i].second,V[i+1].second);
			V.push_back({ID[x],x});
		}
		sort(ALL(V));
		V.erase(unique(ALL(V)),V.end());
		
		ret[0]=N;
		
		vector<int> P;
		FORR(r,V) {
			while(P.size()>1) {
				y = P[P.size()-1];
				if(lca(y,r.second)==y) break;
				x = P[P.size()-2];
				num[x]+=num[y];
				ret[0]-=D[y]-D[x];
				ret[num[y]]+=D[y]-D[x];
				num[y]=0;
				P.pop_back();
			}
			P.push_back(r.second);
		}
		while(P.size()>1) {
			x = P[P.size()-2];
			y = P[P.size()-1];
			num[x]+=num[y];
			ret[0]-=D[y]-D[x];
			ret[num[y]]+=D[y]-D[x];
			num[y]=0;
			P.pop_back();
		}
		
		
		FOR(i,K+1) _P("%d%c",ret[i],(i==K)?'\n':' '), ret[i]=0;
	}
}

まとめ

このLCAで構成された木の構成法およびDFSの方法、どこかで見た気がするけど思い出せず。