[백준 17469] 트리의 색깔과 쿼리

트리에서 쿼리를 처리하는 문제이다.

 

N개의 정점이 있고 각 정점은 1 ~ N까지의 색깔을 가지고 있다.

루트는 항상 1번 정점이고, 트리라고 했으니 간선의 갯수는 항상 N - 1개이다.

 

쿼리가 들어올 때 두 가지 질문을 한다.

1. 두 정점을 끊는다.

2. a 정점에서 갈 수 있는 모든 정점을 보았을 때 색깔 종류의 갯수를 출력한다.

 

이 문제의 조건을 보고 예전에 풀었던 오프라인 쿼리+유니온파인드 관련 문제가 떠올랐다

 

https://www.acmicpc.net/problem/13306

 

문제에서는 간선들이 모두 연결되있고 하나씩 끊어서 처리하라고 했지만, 이 순서대로 하면 구현 난이도도 높고 시간복잡도의 크기가 커진다.

 

그래서 거꾸로 생각을 해 보는 것이다.

모든 정점은 서로 연결되어 있지 않고, 두 정점을 끊는다는 것은 결국 그 전에는 두 정점이 서로 연결되어 있다는 것이다.

 

그래서 입력되는 쿼리를 리스트에 담아서 거꾸로 실행해주면서, 1번 쿼리가 들어오면 두 정점을 연결하고, 2번 쿼리가 들어오면 정답을 입력하는 배열에 담아두었다.

 

각 색깔의 종류들은 서로 중복이 되지 않아야 하기 때문에 TreeSet을 추가로 사용하여야 한다.

TreeSet은 정점의 수 만큼 있어야 하고 두 집합을 Union 할 때 TreeSet의 원소들도 한 쪽으로 합쳐주어야 한다.

하지만 여기서 두 집합의 크기를 고려해야 하는데, 아무렇게나 집합을 합치면 최악의 경우 Union 연산마다 O(N)의 시간복잡도가 들게 된다.

 

그래서 두 집합의 크기를 비교해 크기가 작은 집합에서 크기가 큰 집합으로 Add 하는데 이 때 find()를 통해 찾아진 두 집합 중 크기가 큰 집합이 정답으로 나와야 하니 Disjoint set을 갱신할 때도 find시 항상 큰 집합의 정점 번호가 나오게끔 해 주어야 한다.

---

import java.io.*;
import java.util.*;

public class Main {
    public static int N, Q;
    public static int[] arr, parent, color, ans;
    public static ArrayList<Query> list;
    public static ArrayList<TreeSet<Integer>> treeSetList;
    public static class Query {
        int x, a;
        public Query(int x, int a) {
            this.x = x;
            this.a = a;
        }
    }
    public static void main(String[] args) throws Exception {
        FR fr = new FR();
        N = fr.nI();
        Q = fr.nI();
        arr = new int[N + 1];
        parent = new int[N + 1];
        color = new int[N + 1];
        int qSize = N - 1 + Q;
        ans = new int[Q];
        list = new ArrayList<>(qSize);
        treeSetList = new ArrayList<>(N + 1);
        for(int i = 0; i < N + 1; i++) {
            treeSetList.add(new TreeSet<>());
        }
        for(int i = 1; i <= N; i++) {
            arr[i] = i;
        }
        parent[1] = 1;
        for(int i = 1; i <= N - 1; i++) {
            parent[i + 1] = fr.nI();    //i + 1의 부모는 parent[i + 1]
        }
        for(int i = 1; i <= N; i++) {	// color 작업, 각 번호의 treeSet에 초기값 세팅
            color[i] = fr.nI();
            treeSetList.get(i).add(color[i]);
        }
        for(int i = 0; i < qSize; i++) {	// 오프라인 쿼리
            list.add(new Query(fr.nI(), fr.nI()));
        }
        int idx = Q - 1;
        for(int i = qSize - 1; i >= 0; i--) {
            Query cur = list.get(i);
            if(cur.x == 1) {    // union
                union(cur.a, parent[cur.a]);
            } else {	// 정답 찾기
                int k = find(cur.a);
                ans[idx--] = treeSetList.get(k).size();
            }
        }
        for(int i = 0; i < Q; i++) {
            fr.pl(ans[i]);
        }
        fr.flB();
    }
    public static void union(int a, int b) {
        int x = find(a);
        int y = find(b);
        if(treeSetList.get(x).size() < treeSetList.get(y).size()) {
            treeSetList.get(y).addAll(treeSetList.get(x));
            arr[x] = y;	// size가 큰 번호가 find시 나와야 해서 큰 집합의 번호로 arr을 갱신
        } else {
            treeSetList.get(x).addAll(treeSetList.get(y));
            arr[y] = x;
        }
    }
    public static int find(int target) {
        if(target == arr[target]) {
            return target;
        }
        return arr[target] = find(arr[target]);
    }
}