작은 집합을 큰 집합에 합치기 (Smaller to Larger)

정의

Smaller to Larger (작은 집합을 큰 집합에 합치기) 는 두 집합을 합칠 때 항상 작은 쪽을 큰 쪽에 합치면 전체 O(N log N) 또는 O(N log^2 N) 으로 amortized 비용이 떨어지는 최적화 기법. 트리 서브트리 병합, DSU (Disjoint Set Union), map/set 합병에서 자주 쓰임.

개념 자체는 고전적이지만, “DSU on Tree” (Arpa’s trick) 라는 이름으로 2015년 경 Codeforces 에서 널리 퍼짐.

문제 상황과 동기

트리의 각 서브트리마다 “값의 집합” 을 관리하고, 부모로 올라가며 자식들을 병합.

• naive: 각 노드에서 자식 서브트리를 flat merge. 최악 O(N^2) (일자 트리).
• smaller to larger: 작은 집합을 큰 집합에 합침. 각 원소는 최대 log N 번 이동 → O(N log N) 또는 O(N log^2 N).

핵심 통찰: 한 원소가 합병될 때마다 새 집합 크기가 최소 2배. 따라서 원소당 최대 log N 번 합병.

자주 등장: 트리 서브트리 쿼리 (distinct colors, frequency), DSU 최적화, map/set 병합.

시각화

핵심 아이디어

invariant: 각 원소는 최대 log N 번만 다른 집합으로 복사됨.

merge(A, B):
    if |A| > |B|:
        swap(A, B)
    for x in A:
        B.insert(x)
    return B

• 작은 집합 A 를 큰 집합 B 로 합침.
• A 의 각 원소는 B 로 이동. 다음 합병 시 B 의 크기는 최소 2|A|.
• 원소 하나가 log N 번 이동하면 최종 집합 크기 ≥ 2^(log N) = N.

map/set 인 경우: insert 가 O(log M) 이면 전체 O(N log N × log N) = O(N log^2 N).

알고리즘

기본형 (set 병합)

dfs(v):
    set S_v = {value[v]}
    for each child u:
        S_u = dfs(u)
        if |S_u| > |S_v|:
            swap(S_v, S_u)
        for x in S_u:
            S_v.insert(x)
    answer[v] = process(S_v)
    return S_v

DSU on Tree (모든 서브트리 쿼리)

dfs(v, keep):
    heavy_child = child with max subtree size
    for each light child u:
        dfs(u, false)
    dfs(heavy_child, true)
    for each light child u:
        add(u)
    answer[v] = query()
    if not keep:
        remove(v)

heavy_child 는 재귀 후 그대로 두고, light child 는 재귀 후 삭제. light child 재추가 시 smaller to larger 보장.

구현

// 서브트리의 distinct 값 개수 (smaller to larger, O(N log^2 N))
#include <bits/stdc++.h>
using namespace std;
​
vector<vector<int>> adj;
vector<int> val;
vector<int> ans;
​
set<int> dfs(int v, int p) {
  set<int> S;
  S.insert(val[v]);
  for (int u : adj[v]) {
      if (u == p) continue;
      auto Su = dfs(u, v);
      if (Su.size() > S.size()) swap(S, Su);
      for (int x : Su) S.insert(x);
  }
  ans[v] = (int)S.size();
  return S;
}
​
int main() {
  int n; cin >> n;
  adj.resize(n); val.resize(n); ans.resize(n);
  for (int i = 0; i < n; i++) cin >> val[i];
  for (int i = 0; i < n - 1; i++) {
      int u, v; cin >> u >> v; u--; v--;
      adj[u].push_back(v);
      adj[v].push_back(u);
  }
  dfs(0, -1);
  for (int i = 0; i < n; i++) cout << ans[i] << " ";
  cout << "\n";
}

# 서브트리 distinct 값 (smaller to larger, set 병합)
import sys
sys.setrecursionlimit(200000)
input = sys.stdin.readline
​
adj = []
val = []
ans = []
​
def dfs(v, p):
  S = {val[v]}
  for u in adj[v]:
      if u == p:
          continue
      Su = dfs(u, v)
      if len(Su) > len(S):
          S, Su = Su, S
      S |= Su
  ans[v] = len(S)
  return S
​
n = int(input())
val = list(map(int, input().split()))
adj = [[] for _ in range(n)]
ans = [0] * n
for _ in range(n - 1):
  u, v = map(int, input().split())
  u -= 1; v -= 1
  adj[u].append(v)
  adj[v].append(u)
dfs(0, -1)
print(' '.join(map(str, ans)))

// 서브트리 distinct 값 (smaller to larger, HashSet 병합)
import java.util.*;
import java.io.*;
​
public class Main {
  static List<Integer>[] adj;
  static int[] val, ans;
​
  static HashSet<Integer> dfs(int v, int p) {
      HashSet<Integer> S = new HashSet<>();
      S.add(val[v]);
      for (int u : adj[v]) {
          if (u == p) continue;
          HashSet<Integer> Su = dfs(u, v);
          if (Su.size() > S.size()) {
              HashSet<Integer> tmp = S; S = Su; Su = tmp;
          }
          S.addAll(Su);
      }
      ans[v] = S.size();
      return S;
  }
​
  public static void main(String[] args) throws IOException {
      BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
      int n = Integer.parseInt(br.readLine());
      StringTokenizer st = new StringTokenizer(br.readLine());
      val = new int[n];
      ans = new int[n];
      adj = new ArrayList[n];
      for (int i = 0; i < n; i++) {
          val[i] = Integer.parseInt(st.nextToken());
          adj[i] = new ArrayList<>();
      }
      for (int i = 0; i < n - 1; i++) {
          st = new StringTokenizer(br.readLine());
          int u = Integer.parseInt(st.nextToken()) - 1;
          int v = Integer.parseInt(st.nextToken()) - 1;
          adj[u].add(v);
          adj[v].add(u);
      }
      dfs(0, -1);
      StringBuilder sb = new StringBuilder();
      for (int i = 0; i < n; i++) sb.append(ans[i]).append(' ');
      System.out.println(sb);
  }
}

5
1 2 1 3 2
1 2
1 3
2 4
2 5

3 2 1 1 1

복잡도

각 원소는 최대 log N 번 복사. set 의 insert 가 O(log M) 이면 × log N. 단순 배열 병합은 O(1) insert 라 O(N log N).

변형 / 활용

1. DSU on Tree (Arpa’s trick)

모든 서브트리에서 쿼리. heavy/light decomposition 결합. 각 노드에서 light child 만 재추가 → O(N log N).

2. map 병합 (frequency count)

map<int, int> 를 병합할 때도 동일 원리. 작은 map 을 큰 map 에 insert.

3. persistent data structure

작은 집합을 복사하므로 persistent 버전 (함수형 자료구조) 으로 확장 가능.

함정

1. swap 누락

항상 작은 쪽을 큰 쪽으로 합쳐야 함. swap 안 하면 최악 O(N^2).

2. move semantic 미사용 (C++)

S.insert(Su.begin(), Su.end()) 대신 swap + insert 가 효율적. C++11 move 로 더 빠름.

3. 재귀 깊이 (Python)

Python 은 sys.setrecursionlimit 필수. N=10^5 면 기본 1000 한도 초과.

4. DSU on Tree 의 keep flag 오용

keep=false 일 때 서브트리 정보를 삭제해야 메모리 절약. 안 하면 공간 O(N^2).

BOJ 연습 문제

참고

• 트리 기초
• 깊이 우선 탐색
• 유니온 파인드
• HLD