[BOJ] 17469번 : 트리의 색깔과 쿼리 C++ 풀이

문제를 해석하니 Set + 오프라인 쿼리 + 유니온 파인드 세 가지 기법을 활용하면 무조건 정답을 구할 수 있다는 풀이가 나왔다.

 

우선 모든 간선을 끊어주자. 그리고 쿼리를 뒤에서부터 수행하면서 노드를 하나씩 연결해 주자. 
이때 부모 노드를 찾아줘야 하는데, 부모 노드에서 K개의 색깔을 얻을 수 있다면 해당 집합에 속한 모든 노드는 K개의 색을 얻을 수 있다. 부모 노드를 찾기 위해서는 유니온 파인드 알고리즘을 사용하면 된다. 그리고 색을 중복되지 않게 세줘야 하기 때문에 set으로 관리해 주자.

하지만, 여전히 문제가 있다. 
자식 노드와 부모 노드가 merge 될 때, 자식 노드에 들어있는 색을 부모로 모두 옮겨야 하는데, 이 작업을 수행하는 시간 복잡도는 O(NlogN)이고, 최대 N번 수행되므로 최종 시간복잡도는 O(N^2logN)이 소요된다.

이때 슬쩍 든 생각은 자식과 부모 집합을 비교해서, 현재 가지고 있는 색깔 집합의 크기가 더 큰 노드를 부모로 설정하는 것이었다. 그런데 증명을 못해서 코드를 짜지 않았다.

도저히 방법이 생각나지 않아 알고리즘 분류를 봤고, Small to Large라는 테크닉이 존재하는 것을 알았다. 그리고 이는 위에서 언급한 아이디어를 그대로 가져다 쓰면 되는 테크닉이다.

Small to Large Trick

위의 글을 통해 해당 테크닉을 이해했고, 아주 간단하게 증명할 수 있다.

집합 A, B가 있고 A 집합의 원소 개수는 p, B 집합의 개수는 q라고 하자. 그리고 p >= q이다.
이때 B에 들어있는 원소들을 전부 A로 옮겨주자. 그러면 새로운 집합에 들어있는 원소 개수는 p + q가 되고 p + q >= 2*q는 항상 성립한다. 결국 옮겨지는 원소들은 항상 현재 자신의 집합의 크기보다 2배 이상 큰 집합에 들어가고, 하나의 집합으로 합쳐질 때까지 각 원소들은 logN번 이상 이동되지 않는 것을 알 수 있다.

따라서 이 테크닉을 사용하면 O(NlogN)에 문제가 풀린다.

#include <bits/stdc++.h>
#define MINF 0x7f7f7f7f
#define INF 1000000000
#define MOD 10007
#define NUM 100010
#define X first
#define Y second
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<pii, int> piii;
typedef pair<ll, ll> pll;
typedef pair<double, int> pdi;
int N, Q, par[NUM], r[NUM];
set<int> st[NUM];
vector<pii> v;
vector<int> answer;

int find(int start) {
    if (start == r[start]) return start;
    return r[start] = find(r[start]);
}

void merge(int p, int son) {
    p = find(p), son = find(son);
    if (p == son) return;
    if (st[son].size() < st[p].size()) {
        r[son] = p;
        for (int c : st[son]) st[p].insert(c);
        st[son].clear();
    } else {
        r[p] = son;
        for (int c : st[p]) st[son].insert(c);
        st[p].clear();
    }
    
}

void init() {
    cin >> N >> Q;
    par[1] = 1;
    r[1] = 1;
    for (int i=2; i<=N; i++) {
        cin >> par[i];
        r[i] = i;
    }
    for (int i=1; i<=N; i++) {
        int x;
        cin >> x;
        st[i].insert(x);
    }
    for (int i=1; i<=N+Q-1; i++) {
        int x, y;
        cin >> x >> y;
        v.push_back({x, y});
    }
    reverse(v.begin(), v.end());
    for (auto [x, y] : v) {
        if (x == 1) {
            int p = par[y];
            merge(p, y);
        } else {
            int node = find(y);
            answer.push_back(st[node].size());
        }
    }
    reverse(answer.begin(), answer.end());
    for(int a : answer) cout << a << '\n';
} 

int main(){
    ios_base :: sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    init();
}