yukicoder : No.828 全方位神童数

一応本番中に通したけどだいぶ苦労した。
https://yukicoder.me/problems/no/828

問題

1～Nの番号を持つ頂点からなる木を成すグラフが与えられる。
根が定まっているとき、「自身の頂点番号がどの先祖よりも大きい」ような頂点を神童と呼ぶことにする。
各根に対し、神童となる頂点数を求めよ。

解法

DFSを2回行って解く。
現在神童であるような頂点候補のsetを管理しつつDFSしよう。
DFSの過程で通過した頂点に対し、set中でそれより小さい候補があれば、そのSubTree内でその頂点は神童とならないので一旦それらを取り除く。
ただし取り除いた候補は覚えておき、DFSから抜けるときに戻してやる。
各頂点通過時のset内の要素数の総和を数えればよい。

ただ、これは親子関係にある頂点の神童は列挙できるが、兄弟関係にある頂点同士の処理が行えない。
兄にあたる要素を先にsetに追加した後弟要素を辿ると、弟側の頂点では兄にあたる神童を数え上げられるが逆側はうまくいかない。
よって、2回DFSを行い、2回目は兄弟順を逆にたどることによって兄弟の頂点同士で互いに神童を数え上げる。

祖先の関係にある神童は二重カウントしないよう注意。
以下では、2回のDFSのうち、自身をsetに加えるタイミングを1回目と2回目で子の探索前後で振り分けている。

int N;
ll R[202020];
vector<int> E[202020];
ll S[202020];
int in[202020];
ll mo=1000000007;

set<int> T;

void dfs1(int cur,int pre) {
	vector<int> rem;
	while(T.size()&&*T.begin()<cur) {
		rem.push_back(*T.begin());
		T.erase(T.begin());
	}
	T.insert(cur);
	FORR(e,E[cur]) if(e!=pre) {
		dfs1(e,cur);
		while(T.size()&&*T.begin()<cur) T.erase(T.begin());
	}
	S[cur]+=T.size();
	FORR(r,rem) T.insert(r);
	
}

void dfs2(int cur,int pre) {
	vector<int> rem;
	while(T.size()&&*T.begin()<cur) {
		rem.push_back(*T.begin());
		T.erase(T.begin());
	}
	S[cur]+=T.size();
	FORR(e,E[cur]) if(e!=pre) {
		dfs2(e,cur);
		while(T.size()&&*T.begin()<cur) T.erase(T.begin());
	}
	FORR(r,rem) T.insert(r);
	T.insert(cur);
	
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	
	cin>>N;
	FOR(i,N) cin>>R[i+1];
	
	
	FOR(i,N-1) {
		cin>>x>>y;
		E[x].push_back(y);
		E[y].push_back(x);
	}
	
	dfs1(1,-1);
	FOR(i,N) reverse(ALL(E[i+1]));
	T.clear();
	dfs2(1,-1);
	T.clear();
	
	ll ret=1;
	for(i=1;i<=N;i++) ret=ret*(R[i]+S[i])%mo;
	cout<<ret<<endl;
}

まとめ

DFS訪問順で頂点セットを管理するテクは、２・３回触れたことがあったのでなんとか解けた。