문제 소개
이 문제를 봅시다.
N개의 정점으로 구성된 트리에서, 각각의 정점들은 색깔을 하나씩 갖고 있습니다.
쿼리가 두 가지 주어지는데
- 1번 쿼리는 정점 a와 a의 부모 정점을 잇는 간선을 끊는 쿼리
- 2번 쿼리는 정점 a에서 갈 수 있는 정점들을 봤을 때, 서로 다른 색깔의 종류를 구하는 쿼리 가 꽤 많이 주어집니다.
이 문제는 2016KOI 중등3번문제(풀이)에 색깔 개념을 추가한 문제입니다.
13306번처럼 쿼리의 순서를 뒤집어서 간선을 끊는 쿼리 대신 간선을 연결해주는 쿼리 로 바꾸면 1번 쿼리는 Union-Find로 해결할 수 있습니다.
2번 쿼리는 서로 다른 색깔의 종류를 구하는 쿼리입니다. 어떤 정점에서 갈 수 있는 정점들의 색깔을 모두 set에 넣은 다음에 사이즈를 출력하면 쉽게 해결할 수 있을 것 같습니다.
모든 정점은 처음에 분리된 상태에서 시작하기 때문에, 각 정점마다 set을 하나씩 만들 것입니다. 처음에는 각 정점의 색깔을 set에 넣어준 상태로 시작합니다.
두 정점 u, v를 잇는 간선을 추가할 때 우리는 Union-Find에서 union 연산을 해주게 됩니다. 이때 두 정점의 set도 합쳐주는 작업을 같이 해준다면, 각 정점에서 갈 수 있는 정점의 색깔을 관리해줄 수 있습니다.
Naive Solution
정점 u를 담당하는 set과 정점 v를 담당하는 set을 합치는 것은 쉽게 짤 수 있습니다.
1 | |
최악의 경우를 살펴봅시다.
처음에 크기가 1인 두 set을 합치고(원소 1번 이동)
크기가 2인 set의 원소를 크기가 1인 또 다른 set에 넣어주고(원소 2번 이동)
크기가 3인 set의 원소를 크기가 1인 또 다른 set에 넣어주고(원소 3번 이동)
…
set을 잘못된 방법으로 합쳐주게 된다면, 최악의 경우에는 원소가 총 O(N2)번 이동할 수도 있습니다.
위 예시를 보면서 작은 set의 원소를 큰 set에 옮겨주면 어느정도 커팅이 될 것 같다는 생각이 들 수 있습니다.
사실 약간의 커팅이 아닌, O(N2)을 O(N log N)으로 바꿀 수 있는 강력한 방법입니다.
small to Large
이 글의 주제입니다.
작은 set의 원소를 큰 set에 옮기면 어떤 일이 발생하는 지 알아봅시다.
처음에 N개의 원소들은 모두 크기가 1인 set에 들어있습니다. 그리고 모든 set을 다 합쳐주게 되면 크기가 N인 set에 들어가게 됩니다.
작은 set에 있는 원소를 큰 set으로 옮겨주면, 옮겨지는 원소가 속한 set의 크기는 항상 2배 이상 증가 합니다.
옮겨지는 원소가 속한 set의 크기는 2배 이상 증가하기 때문에, 크기가 1인 set에 있는 원소가 크기가 N인 set으로 이동할 때, 최대 O(log N)번만 이동하게 됩니다.
결국 각 원소는 최대 O(log N)번 이동하고, 원소의 이동은 최대 O(N log N)번만 발생하게 됩니다.
set의 insert연산은 O(log N)이므로 O(N log2 N)에 문제를 해결할 수 있습니다.
Heavy Light Decomposition에도 이 원리가 활용되고, Union-Find의 Union-By-Rank나 Union-By-Size에도 이 원리가 사용되므로 알아두면 유용한 것 같습니다.