문제 21
집합 합치기
처음에 $1$번부터 $N$번까지의 원소가 각각 다른 집합에 속해 있다.
$Q$개의 쿼리를 순서대로 처리하라.
쿼리는 다음 중 하나이다.
merge a b: $a$가 속한 집합과 $b$가 속한 집합을 합친다.same a b: $a$와 $b$가 같은 집합에 속하면1, 아니면0을 출력한다.size a: $a$가 속한 집합의 크기를 출력한다.
입력
첫 번째 줄에 두 정수 $N$, $Q$가 공백으로 구분되어 주어진다. $(1 \le N, Q \le 10^6)$
두 번째 줄부터 $Q$개의 줄에 걸쳐 쿼리가 하나씩 주어진다.
쿼리에 포함되는 원소 번호 $a$, $b$의 범위는 다음과 같다. $(1 \le a, b \le N)$
출력
same, size 쿼리마다 결과를 한 줄에 하나씩 출력한다.
예제 입력 1
8 10 same 1 2 merge 1 2 same 1 2 merge 2 3 size 1 size 3 merge 4 5 merge 3 5 same 1 4 size 2
예제 출력 1
0 1 3 3 1 5
공식 해설
각 합치기 때 원소들의 소속을 전부 바꾸면 많은 원소를 반복 방문하게 된다. $N,Q\le10^6$이므로 집합마다 대표를 두어 필요한 연결만 바꾸자.
집합이 합쳐지기만 하므로 분리 집합 자료구조(DSU)를 사용한다. 각 원소에 부모를 저장하고, 부모를 따라 도착하는 루트를 그 집합의 대표로 삼는다. 처음에는 각 원소가 자기 자신을 대표로 가지며 집합 크기는 $1$이다.
same a b는 두 원소의 대표가 같은지 비교한다. size a는 $a$의 대표에 저장된 크기를 출력한다.
merge a b에서는 두 대표를 찾는다. 같다면 이미 한 집합이므로 아무것도 하지 않는다. 다르다면 작은 집합의 루트를 큰 집합의 루트에 연결하고 크기를 더한다. 집합 크기는 대표에서만 관리한다.
대표를 찾은 뒤 지나온 원소들의 부모를 그 대표로 바꾸는 경로 압축을 함께 적용한다. 크기 기준 합치기와 경로 압축을 사용하면 전체 시간복잡도는 $O((N+Q)\alpha(N))$, 공간복잡도는 $O(N)$이다. 여기서 $\alpha$는 역 아커만 함수이다.
구현 · 언어별 풀이
C++ 구현
vector<int>에 부모와 집합 크기를 저장한다. 처음 부모는 자기 자신으로, 크기는 $1$로 초기화한다. 대표를 찾는 함수에 경로 압축을 적용하고, 합칠 때는 크기를 비교해 작은 루트를 큰 루트에 붙인다. 크기 조회도 원소 자신이 아니라 대표의 값을 읽는다.
C 구현
parent와 size를 힙에 할당한 int 배열로 저장한다. find는 반복문으로 루트를 찾고 경로를 압축하며, 작은 트리를 큰 트리에 붙인다.
Python 구현
parent와 size에는 array('i')를 사용할 수 있다. 백만 개의 Python 정수 객체를 여러 배열에 중복 보관하는 비용을 줄인다. find는 반복문과 경로 압축으로 구현한다.
Java 구현
int[] parent와 size를 사용한다. find는 반복문으로 경로를 압축하고 size가 작은 루트를 큰 루트에 붙인다.
Rust 구현
Vec<usize>로 parent와 size를 저장한다. find는 반복문으로 루트를 찾고 경로를 압축한다.
JavaScript 구현
부모와 집합 크기를 각각 Int32Array(n + 1)로 저장한다. 최대 백만 개의 번호·크기는 모두 32비트 범위에 들어간다. 초기 부모는 자기 번호, 크기는 $1$이다.
find는 반복문으로 루트까지 올라가 경로를 압축하고, 합칠 때는 작은 집합의 대표를 큰 대표 아래에 붙인다. size a는 sz[find(a)]이다. 백만 개의 질의는 바이트 입력을 순차로 읽어 처리하면 전체 토큰 문자열을 장기간 저장하지 않아도 된다.