DSU & Kruskal¶
In one sentence
A DSU (Disjoint Set Union, also called Union-Find) keeps track of groups that only ever merge. It answers "are \(x\) and \(y\) in the same group?" almost instantly. Kruskal's algorithm uses a DSU to build the cheapest network that connects everything.
1. What problem does it solve?¶
Example — CF 1167C · News Distribution
There are \(n\) users and \(m\) friend groups. A user who learns some news tells every group they belong to, those people tell their groups, and so on.
For each user \(x\): if \(x\) starts spreading the news, how many users end up knowing it?
Limits: \(n, m \le 5 \cdot 10^5\), total group sizes \(\le 5 \cdot 10^5\).
Input Output
7 5
3 2 5 4 4 4 1 4 4 2 2
0
2 1 2
1 1
2 6 7
Being in the same group is contagious. If \(2, 5, 4\) share a group and \(1, 2\) share a group, then \(1, 2, 4, 5\) all hear the same news. So the users split into components, and the answer for \(x\) is the size of \(x\)'s component:
graph LR
1 --- 2
2 --- 5
2 --- 4
3
6 --- 7Components: \(\{1, 2, 4, 5\}\) has size 4, \(\{3\}\) has size 1, and \(\{6, 7\}\) has size 2. That matches the output 4 4 1 4 4 2 2.
We need two operations, many times:
| operation | meaning | here |
|---|---|---|
union(u, v) |
merge the groups of \(u\) and \(v\) | a friend group links its members |
find(x) |
name of \(x\)'s group (its root) | whose component is \(x\) in? |
2. How it works¶
2.1 Groups as trees¶
Every group is stored as a tree. Each element points to a parent, and the root points to itself. The root is the group's name.
- Start:
parent[x] = x. Every element is its own one-element group. find(x): followparentlinks until you reach an element that points to itself.
union(u, v): find both roots. If they differ, hang one root under the other. It is one assignment.
Naive union: root \(1\) is hung under root \(6\). With union by size (§2.4) it would be the other way round, because \(1\)'s tree has 5 elements and \(6\)'s has 3.
Two elements are in the same group exactly when find(u) == find(v).
2.2 Why the naive version can be slow¶
If we always hang the second root under the first, a bad order of unions builds a chain:
union(2,1), union(3,2), union(4,3), ... 1 <- 2 <- 3 <- 4 <- ... <- n
Now find(n) walks \(n\) steps, and \(n\) such calls cost \(O(n^2)\). There are two fixes, and we use both.
2.3 Fix 1 — path compression¶
While walking up to the root, reconnect every visited element directly to the root. The next find on any of them takes one step.
Each line goes up to the parent. The red lines are the path find(4) walked, and the links it rewired.
This changes the shape of the tree, but not who is in which group, because every element still has the same root.
2.4 Fix 2 — union by size¶
When merging, hang the smaller tree under the larger one, and keep sz[root] = number of elements in the group.
Claim: with union by size, every tree has height at most \(\log_2 n\).
Why: an element's depth increases by \(1\) only when its tree is hung under a tree that is at least as big. After that merge its group is at least twice as large as before. A group can double at most \(\log_2 n\) times before it contains all \(n\) elements, so the depth is at most \(\log_2 n\).
| version | find cost |
|---|---|
| naive | \(O(n)\) worst case |
| union by size | \(O(\log n)\) |
| path compression + union by size | \(O(\alpha(n))\), where \(\alpha(n) \le 4\) for any real input: effectively constant |
2.5 Kruskal's algorithm (Minimum Spanning Tree)¶
Problem. A connected, undirected graph has \(n\) vertices and weighted edges. Choose edges so that all vertices are connected and the total weight is as small as possible. The result is a tree with \(n - 1\) edges, called the minimum spanning tree (MST).
Algorithm.
- Sort the edges by weight, smallest first.
- For each edge \((u, v)\) in that order: if \(u\) and \(v\) are not already connected, take the edge and
union(u, v). Otherwise skip it, because it would create a cycle. - Stop after taking \(n - 1\) edges.
The DSU answers "already connected?" in step 2.
Example — 1095F · Make It Connected (6-DSU · I, rating 1900):
Problem
\(n\) vertices with numbers \(a_1, \dots, a_n\) and no edges. Adding edge \((x, y)\) costs \(a_x + a_y\). There are also \(m\) special offers \((x, y, w)\): you may add edge \((x, y)\) for \(w\) instead. Find the minimum cost to make the graph connected.
Limits: \(n, m \le 2 \cdot 10^5\), \(a_i, w \le 10^{12}\).
Input Output
3 2
1 3 3
2 3 5
2 1 1 5
4 0
1 3 3 7 16
5 4
1 2 3 4 5
1 2 8
1 3 10
1 4 7
1 5 15 18
This is an MST, but on a complete graph with \(\frac{n(n-1)}{2}\) ordinary edges, which is far too many to sort. The trick: let \(v_0\) be the vertex with the smallest \(a\). Any ordinary edge \((x, y)\) with \(x, y \ne v_0\) costs \(a_x + a_y \ge a_{v_0} + a_y\). So if an MST used \((x, y)\), removing it splits the tree into two parts; \(v_0\) is in one of them, and the edge from \(v_0\) to whichever of \(x, y\) is in the other part reconnects them for no more cost. Therefore we only need the \(n - 1\) edges \((v_0, i)\) plus the \(m\) offers: about \(4 \cdot 10^5\) edges.
Trace on the first sample: \(a = [1, 3, 3]\), so \(v_0 = 1\). Edges: \((1, 2)\) costs \(4\), \((1, 3)\) costs \(4\), and the offers \((2, 3)\) for \(5\) and \((2, 1)\) for \(1\).
graph LR
1 ---|"offer 1"| 2
1 ---|"a1+a3 = 4"| 3
1 ---|"a1+a2 = 4"| 2
2 ---|"offer 5"| 3
linkStyle 0,1 stroke:#4caf50,stroke-width:4px| edge (sorted) | weight | same component already? | action | total |
|---|---|---|---|---|
| 2 – 1 (offer) | 1 | no | take, union(2, 1) |
1 |
| 1 – 3 | 4 | no | take, union(1, 3); now \(2 = n - 1\) edges, stop |
5 |
| 1 – 2 | 4 | (yes) | — | |
| 2 – 3 (offer) | 5 | (yes) | — |
The green edges form the MST, with total weight \(5\).
Why greedy is correct (the exchange argument). Suppose Kruskal takes edge \(e\), but some optimal MST \(T\) does not contain it. Adding \(e\) to \(T\) creates a cycle. That cycle must contain another edge \(f\) that crosses between the two components \(e\) was joining. Kruskal had not yet taken \(f\) at that moment, so \(f\) comes later in sorted order and \(w(f) \ge w(e)\). Swap \(f\) out and \(e\) in: the result is still a spanning tree, and its weight is not larger. Repeating this turns any optimal tree into Kruskal's tree without increasing the weight, so Kruskal's tree is optimal.
Cost: sorting takes \(O(m \log m)\), and the DSU work is almost linear. So \(O(m \log m)\) overall.
3. Templates¶
These are the templates we use in class. Use them exactly as written.
Naive DSU¶
parent = [i for i in range(n + 1)]
def find(x):
while parent[x] != x:
x = parent[x]
return x
def union(u, v):
pu = find(u)
pv = find(v)
if pu != pv:
parent[pv] = pu
DSU with Path Compression¶
def find(x):
if parent[x] != x:
parent[x] = find(parent[x])
return parent[x]
Only find changes.
DSU with Path Compression + Union by Size¶
parent = [i for i in range(n + 1)]
sz = [1 for _ in range(n + 1)]
def find(x):
if parent[x] != x:
parent[x] = find(parent[x])
return parent[x]
def union(u, v):
pu = find(u)
pv = find(v)
if pu != pv:
if sz[pu] < sz[pv]:
pu, pv = pv, pu
parent[pv] = pu
sz[pu] += sz[pv]
Kruskal¶
def solve():
n, m = map(int, input().split())
edges = []
for _ in range(m):
u, v, w = map(int, input().split())
edges.append((w, u, v))
edges.sort()
parent = [i for i in range(n + 1)]
sz = [1 for _ in range(n + 1)]
def find(x):
if parent[x] != x:
parent[x] = find(parent[x])
return parent[x]
def union(u, v):
pu = find(u)
pv = find(v)
if pu == pv:
return False
if sz[pu] < sz[pv]:
pu, pv = pv, pu
parent[pv] = pu
sz[pu] += sz[pv]
return True
ans = 0
cnt = 0
for w, u, v in edges:
if union(u, v):
ans += w
cnt += 1
if cnt == n - 1:
break
print(ans)
t = int(input())
for _ in range(t):
solve()
4. Worked solution — News Distribution¶
Union every member of a group with the group's first member. The answer for \(x\) is sz[find(x)].
import sys
input = sys.stdin.readline
def solve():
n, m = map(int, input().split())
parent = [i for i in range(n + 1)]
sz = [1 for _ in range(n + 1)]
def find(x):
if parent[x] != x:
parent[x] = find(parent[x])
return parent[x]
def union(u, v):
pu = find(u)
pv = find(v)
if pu != pv:
if sz[pu] < sz[pv]:
pu, pv = pv, pu
parent[pv] = pu
sz[pu] += sz[pv]
for _ in range(m):
g = list(map(int, input().split()))
k = g[0]
for i in range(2, k + 1):
union(g[1], g[i])
for x in range(1, n + 1):
print(sz[find(x)], end=" ")
print()
solve()
Why this passes and a naive version gets TLE
With \(5 \cdot 10^5\) users, unions in a bad order can build long chains. Union by size keeps every tree shallow, and path compression flattens it further. The recursion in find is safe here because union by size keeps the depth at most \(\log_2 n \approx 19\).
5. Worked solution — Make It Connected¶
This is the Kruskal template on one test case. The edge list is the \(m\) offers plus the \(n - 1\) edges from the smallest vertex \(v_0\) (see §2.5 for why no other ordinary edge is needed).
import sys
input = sys.stdin.readline
def solve():
n, m = map(int, input().split())
a = [0] + list(map(int, input().split()))
edges = []
for _ in range(m):
u, v, w = map(int, input().split())
edges.append((w, u, v))
v0 = 1
for i in range(2, n + 1):
if a[i] < a[v0]:
v0 = i
for i in range(1, n + 1):
if i != v0:
edges.append((a[v0] + a[i], v0, i))
edges.sort()
parent = [i for i in range(n + 1)]
sz = [1 for _ in range(n + 1)]
def find(x):
if parent[x] != x:
parent[x] = find(parent[x])
return parent[x]
def union(u, v):
pu = find(u)
pv = find(v)
if pu == pv:
return False
if sz[pu] < sz[pv]:
pu, pv = pv, pu
parent[pv] = pu
sz[pu] += sz[pv]
return True
ans = 0
cnt = 0
for w, u, v in edges:
if union(u, v):
ans += w
cnt += 1
if cnt == n - 1:
break
print(ans)
solve()
6. Common mistakes¶
Linking the elements instead of the roots
parent[v] = u is wrong. It rewires only \(v\) and can break other trees. Always link roots: parent[pv] = pu.
Updating sz on a non-root
sz is only meaningful at roots. Read the size with sz[find(x)], not sz[x].
Comparing parent[u] == parent[v]
Two elements can be in the same group with different parents. Compare find(u) == find(v).
Off-by-one on array size
Vertices are numbered \(1..n\), so allocate n + 1, as in the templates.
Deep recursion in find
Recursive find without union by size can reach depth \(n\) and crash with RecursionError. Use the full template (compression and size). If you must use the naive union, use the iterative find from the naive template.
Kruskal: forgetting to sort, or sorting by the wrong key
Store edges as (w, u, v) so edges.sort() sorts by weight first.
7. Practice¶
| Problem | Set | Rating | Idea |
|---|---|---|---|
| 217A · Ice Skating | 6-DSU · E | 1200 | components = answer + 1 |
| 1167C · News Distribution | 6-DSU · B | 1400 | component sizes (this page) |
| 277A · Learning Languages | 6-DSU · A | 1400 | union employees through languages |
| 1559D1 · Mocha and Diana (Easy) | 6-DSU · H | 1400 | two DSUs at once |
| 1702E · Split Into Two Sets | 6-DSU · D, 5-Graph · D | 1600 | components + odd cycles |
| 25D · Roads not only in Berland | 6-DSU · G | 1900 | redundant edges vs. components |
| 1095F · Make It Connected | 6-DSU · I | 1900 | Kruskal with extra "cheapest vertex" edges |
| 1245D · Shichikuji and Power Grid | 6-DSU · C | 1900 | Kruskal with a virtual power-station vertex |
Make It Connected can always be connected (every pair of vertices has an ordinary edge), so the solution never needs an "impossible" case.
Credits & licenses
- DSU figures in §2.1 (the forest,
find, andunion): copied unchanged from OI Wiki — Disjoint Set Union by HeRaNO, JuicyMio, Xeonacid, sailordiary, ouuan, Pig-Eat-Earth and other OI Wiki contributors, licensed CC BY-SA 4.0. The path-compression figure is ours, because the original has a Chinese label. - Choice of example problems and parts of the practice list: follow USACO Guide — Disjoint Set Union by Benjamin Qi, Andrew Wang and Nathan Gong, and USACO Guide — Minimum Spanning Trees by Benjamin Qi, Andrew Wang and Kuan-Hung Chen, licensed CC BY-NC-SA 4.0.
- Everything else (text, proofs, the Kruskal trace, the worked solutions) is ours. The templates in §3 are from our lessons.