백준 25549 - 트리의 MEX (C++)

문제

• 문제 링크

  
   BOJ 25549 - 트리의 MEX 
  

• 문제 요약

$N$개의 정점으로 이루어진 트리가 주어진다.
$mex(i)$의 정의가 아래와 같을 때, $mex(1), mex(2), ... mex(N)$을 출력하자.



$mex(i)$ : $i$번 정점을 루트로 하는 서브 트리의 정점에 적혀 있지 않은 수 중에서 가장 작은 음이 아닌 정수.


• 제한

TL : $2$ sec, ML : $1024$ MB




$1 ≤ N ≤ 200,000$

알고리즘 분류

• 자료 구조 (data_structures)
• 트리 (trees)
• 트리를 사용한 집합과 맵 (tree _ set / map)
• 작은 집합에서 큰 집합으로 합치는 테크닉 (smaller_to_larger)

풀이

어떤 정점의 서브트리에 포함되는 $v_i$들은 집합을 통해 쉽게 관리할 수 있다.

나이브하게 볼 때, 존재하는 모든 서브 트리에 대한 집합을 보고 $0, 1, ... N-1$ 로 늘려보며 답을 찾아줄 수 있다.
이를 그대로 구현한다면 약 $O(N^2logN)$의 복잡도로 문제를 풀테고 시간 초과를 피할 수 없다.

$mex(i)$를 구할 때, $i$의 자식 노드들에 기록된 $mex$ 값을 재활용할 수 없을까 $?$

자식 노드들의 집합을 합칠 때, 효율적으로 합칠 수 없을까 $?$

두번째 대목에서 $Small$_$To$_$Large$ $Trick$ 을 적용한다면, 요소 별 이동 횟수를 최악 $N$번에서 $logN$번으로 줄일 수 있게 된다. 즉, $O(Nlog^2N)$의 복잡도로 문제를 풀 수 있다.


이제 첫번째 대목을 보자.

매번 $mex(i)$의 시작값을 $0$부터 차례차례 올리는 식으로 작성하기보다, $i$의 가장 규모가 큰 자식 노드에 적힌 $mex$ 값부터 올려보자.

그러면 좀 더 효울적으로 $mex(i)$를 특정할 수 있을 것이다.

전체 코드

#include<bits/stdc++.h>
#define N 200001
using namespace std;
 
vector <int> Gr[N], ans(N);
set <int> mex[N];
int n, r;
 
void Input()
{
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n;
 
    for (int i(1); i <= n; i++)
    {
        int x; cin >> x;
        if (!~x) r = i;
        else Gr[x].push_back(i);
    }
    for (int i(1); i <= n; i++)
    {
        int x; cin >> x;
        mex[i].insert(x);
    }
}
void getmex(int now)
{
    if (!Gr[now].size())
    {
        ans[now] = !*mex[now].begin();
        return;
    }
 
    int i(Gr[now][0]);
    for (int& next : Gr[now])
    {
        getmex(next);
        if (mex[i].size() < mex[next].size()) i = next;
    }
 
    mex[i].insert(*mex[now].begin());
    for (int& next : Gr[now])
        if (next ^ i)
            for (auto& x : mex[next])
                mex[i].insert(x);
 
    swap(mex[i], mex[now]);
 
    int val(ans[i]);
    for (auto it(mex[now].find(val)); it != mex[now].end() && *it == val; it++)
        val++;
 
    ans[now] = val;
}
int main()
{
    Input();
    getmex(r);
    for (int i(1); i <= n; cout << ans[i++] << '\n');
}

comment

끝