yukicoder : No.2618 除霊

やることが決まってもちょっと手間のかかる問題。
https://yukicoder.me/problems/no/2618

問題

木を成す無向グラフが与えられる。
うち指定されたいくつかの点にはお化けがいる。
お化けが魔法を放つと、距離1以内の点が攻撃される。

今点を1つ指定し、そこから距離1以内の点にいるお化けを除霊できるとする。
各点を除霊したケースそれぞれにおいて、その後お化けが攻撃可能な頂点数を求めよ。

解法

まず複数始点のBFSで、各点が何点から攻撃されそうかを求めよう。
次に、本来攻撃される各点に対し、どこを除霊したら攻撃を防げるかを求める。
その範囲は、やはりある点から距離1以内の点となる。
そこで頂点をBFS訪問順に並べ替えておくと、定数個の区間で表現できるのであとは累積和で計算可能である。

int N,M;
vector<int> E[202020];
int D[202020];
int C[202020];
vector<int> B[202020];
int L[202020];
int P[202020],S[202020],T[202020];

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);
	}
	cin>>M;
	FOR(i,M) {
		cin>>x;
		C[x-1]=1;
		FORR(e,E[x-1]) B[e].push_back(x-1);
	}
	int sum=0;
	
	
	MINUS(L);
	queue<int> Q;
	int id=0;
	Q.push(0);
	while(Q.size()) {
		x=Q.front();
		Q.pop();
		L[x]=id++;
		FORR(e,E[x]) if(L[e]==-1) Q.push(e);
	}
	
	FOR(i,N) {
		S[i]=1<<20,T[i]=-1;
		FORR(e,E[i]) {
			if(L[e]<L[i]) P[i]=L[e];
			else {
				S[i]=min(S[i],L[e]);
				T[i]=max(T[i],L[e]);
			}
		}
	}
	
	FOR(i,N) {
		if(C[i]||B[i].size()) {
			sum++;
		}
		else {
			continue;
		}
		if(B[i].size()>=2) {
			D[L[i]]++;
			D[L[i]+1]--;
		}
		else if(B[i].empty()) {
			if(i) {
				D[P[i]]++;
				D[P[i]+1]--;
			}
			D[L[i]]++;
			D[L[i]+1]--;
			if(S[i]<=N) {
				D[S[i]]++;
				D[T[i]+1]--;
			}
		}
		else {
			x=B[i][0];
			if(C[i]) {
				D[L[i]]++;
				D[L[i]+1]--;
				D[L[x]]++;
				D[L[x]+1]--;
			}
			else {
				if(x) {
					D[P[x]]++;
					D[P[x]+1]--;
				}
				D[L[x]]++;
				D[L[x]+1]--;
				if(S[x]<1<<20) {
					D[S[x]]++;
					D[T[x]+1]--;
				}
			}
		}
	}
	FOR(i,N) D[i+1]+=D[i];
	FOR(i,N) cout<<sum-D[L[i]]<<endl;
	
	
}

まとめ

最初BIT使おうとしてしまった。