BOJ 25549 트리의 MEX

문제 링크입니다.

트리가 직선형이라면

$i = 2 \cdots N$에 대해 $p_i = i-1$이라고 합시다. 즉 $1$번 정점의 자식으로 $2$번 정점이, $2$번 정점의 자식으로 $3$번 정점이 … $N-1$번 정점의 자식으로 $N$번 정점이 오게끔 입력이 주어졌다고 합시다.

$i$번 정점을 루트로 하는 서브트리는 구간 $[i, N]$로 나타납니다. $N$번 정점부터 역으로 올라가면서 답을 구해줘 봅시다.

집합에 새로운 수가 들어온다고 해서 mex가 감소하는 일은 없습니다. 그러니 마치 투 포인터를 하듯이 mex를 관리해 줄 수 있습니다. 정점을 거슬러 올라가면서, 집합에 이번 정점의 값을 넣고, 집합에 없는 값이 나올 때까지 mex 값을 올리면서 탐색을 진행해 주면 됩니다. mex 값은 $N$을 넘지 않기 때문에, mex의 값을 변경하는 작업은 최대 $N$번 일어납니다. 시간 복잡도에는 별 영향을 주지 못하겠네요. 총 시간 복잡도는 $\mathcal{O}(N \log N)$입니다.

일반적인 트리의 경우

트리를 직선형 여러 개가 한 정점에 들러붙어 있는 느낌으로 생각해 줍시다. 직선형의 경우와 마찬가지로 각 자식 정점에 대한 집합을 구해 준 뒤에 그걸 부모 정점의 값과 합친 커다란 집합을 만들어 주어야 합니다. Smaller-to-Larger 테크닉을 활용하면 이를 $\mathcal{O}(N \log ^2 N)$에 수행할 수 있습니다.

각 정점마다 mex 값을 저장하면서, 부모 정점의 mex 값을 구할 때 활용해 줍시다. 직선형 때와 마찬가지로 부모 정점의 mex 값이 자식 정점 이상임을 이용해 투 포인터스럽게 구현해 주면 총 시간복잡도 $\mathcal{O}(N \log ^2 N)$에 문제를 해결할 수 있습니다.

코드

#include<bits/stdc++.h>
using namespace std;
using ll=long long;
vector<vector<int>> e;
vector<set<int>> s;
int ord[202020];
int ans[202020];
void dfs(int now){
	int mex = 0;
	for(int nxt : e[now]){
		dfs(nxt);
		mex = max(mex, ans[nxt]);
		if(s[ord[now]].size() < s[ord[nxt]].size()) swap(ord[now], ord[nxt]);
		for(auto x : s[ord[nxt]]) s[ord[now]].insert(x);
	}
	for(int i = mex; ; i++) if(s[ord[now]].find(i) == s[ord[now]].end()){
		ans[now] = i; return;
	}
}
int main(){
	cin.tie(0);cout.tie(0);ios::sync_with_stdio(0);
	int n; cin >> n; e.resize(n+1); s.resize(n+1);
	int start = 0;
	for(int i = 1; i <= n; i++){
		ord[i] = i;
		int par; cin >> par;
		if(par > 0) e[par].push_back(i);
		else start = i;
	}
	for(int i = 1; i <= n; i++){
		int x; cin >> x;
		s[i].insert(x);
	}
	dfs(start);
	for(int i = 1; i <= n; i++) cout << ans[i] << '\n';
}