Структура данных Union-Find: как объединять множества за O(1)

Зачем вам Union-Find?

Представьте: вы решаете задачу, где нужно объединять группы друзей по мере появления новых связей. Первая мысль — хранить список списков и при каждом объединении сканировать всё. Это работает, но на большом тесте программа падает с Time Limit Exceeded. Знакомо? Union-Find (или DSU) решает эту проблему элегантно и быстро.

Если вы когда-нибудь застревали, пытаясь отслеживать связные компоненты при добавлении рёбер, вы по адресу. Давайте превратим фрустрацию в суперсилу.

Что такое Union-Find на самом деле

Представьте n изолированных островов. Со временем между ними строят мосты. Нужно быстро отвечать на два вопроса:

• Соединены ли два острова (напрямую или через другие)?
• Объединить группы, содержащие два острова.

Наивное решение сканирует все острова каждый раз — сложность O(n²). Union-Find обходит это, храня для каждого элемента указатель на его «родителя» — представителя множества. Магия в двух идеях:

1. Сжатие пути — при поиске корня мы перенаправляем все узлы на пути прямо к корню. Последующие поиски становятся мгновенными.
2. Объединение по рангу (или размеру) — при слиянии двух деревьев подвешиваем меньшее под корень большего. Это сохраняет дерево неглубоким.

Амортизированная стоимость последовательности из m операций find и union на n элементах — O(m α(n)), где α — обратная функция Аккермана. Она растёт так медленно, что для любых разумных входных данных α(n) ≤ 5. В контексте собеседований можно считать каждую операцию O(1).

Почему это работает (интуиция)

Массив parent кодирует лес, где корень каждого дерева — «метка» множества. При объединении мы просто делаем один корень родителем другого. Сжатие пути не меняет логическую группировку, а лишь переподключает указатели, чтобы будущие вызовы find пропускали лишние шаги.

Представьте очередь людей, передающих ведро по цепочке. Без сжатия ведро проходит через каждого. Со сжатием после первой передачи все получают телепорт к началу — ведро прибывает мгновенно.

Код: класс Union-Find на Python

Ниже — компактная, готовая к продакшену реализация с комментариями.

class UnionFind:
    def __init__(self, n: int):
        self.parent = list(range(n))  # каждый элемент — свой родитель
        self.rank = [0] * n           # приблизительная глубина дерева

    def find(self, x: int) -> int:
        # Поиск корня со сжатием пути
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int) -> bool:
        # Объединение по рангу; возвращает True, если произошло слияние
        xr, yr = self.find(x), self.find(y)
        if xr == yr:
            return False  # уже в одном множестве
        if self.rank[xr] < self.rank[yr]:
            self.parent[xr] = yr
        elif self.rank[xr] > self.rank[yr]:
            self.parent[yr] = xr
        else:
            self.parent[yr] = xr
            self.rank[xr] += 1
        return True

Частые ошибки (они же «босс-файты»)

• Забыть сжатие пути — если find реализован простым циклом без обновления родителей, каждая операция может деградировать до O(log n) или хуже.
• Объединение без ранга/размера — произвольное подвешивание может создать длинные цепочки, превращая структуру в связный список в худшем случае.

Реальные задачи с собеседований

Задача 1: Количество островов (LeetCode 200)

Дана двумерная сетка из '1' (земля) и '0' (вода). Нужно подсчитать количество островов. Соседние по горизонтали или вертикали клетки принадлежат одному острову.

Решение через Union-Find: обрабатываем каждую клетку как узел. Сканируем сетку; для каждой земляной клетки объединяем её с верхним и левым соседом (если они тоже земля). В конце количество различных корней среди земляных клеток равно числу островов.

def numIslands(grid):
    if not grid:
        return 0
    m, n = len(grid), len(grid[0])
    uf = UnionFind(m * n)
    count = sum(cell == '1' for row in grid for cell in row)
    for i in range(m):
        for j in range(n):
            if grid[i][j] == '1':
                idx = i * n + j
                if i > 0 and grid[i-1][j] == '1':
                    if uf.union(idx, (i-1)*n + j):
                        count -= 1
                if j > 0 and grid[i][j-1] == '1':
                    if uf.union(idx, i*n + (j-1)):
                        count -= 1
    return count

Почему это круто: вместо запуска DFS/BFS для каждой новой клетки (что может приводить к повторным обходам) каждая операция union/find практически O(1). Общая сложность O(m·n) с минимальными константами.

Задача 2: Лишнее ребро (LeetCode 684)

В дереве с n вершинами (пронумерованы от 1 до n) добавлено одно лишнее ребро, создающее единственный цикл. Нужно вернуть это ребро.

Решение через Union-Find: проходим по рёбрам; для каждого ребра (u, v), если find(u) == find(v), то добавление ребра создаст цикл — это лишнее ребро. Иначе объединяем множества.

def findRedundantConnection(edges):
    uf = UnionFind(len(edges) + 1)  # вершины нумеруются с 1
    for u, v in edges:
        if not uf.union(u, v):
            return [u, v]

Алгоритм работает за O(E α(V)) ≈ O(E) времени и O(V) памяти — оптимально для этой задачи.

Почему это важно

Освоение Union-Find — это не просто подготовка к LeetCode. Это ментальный инструмент для любых сценариев, где связи эволюционируют во времени: динамические сети, перколяция, обработка изображений, фича социальной сети «вы в одном сообществе?».

Когда вы можете ответить «связаны ли эти два элемента?» почти за константу, вы освобождаетесь от квадратичных ловушек. Задачи начинают выглядеть как множества и слияния, и кажущееся сложным становится простым.

Ваш ход

Выберите задачу, которая у вас не шла, связанную с группировкой или связностью — например, «Дружеские круги» (LeetCode 547) или «Слияние аккаунтов» (LeetCode 721). Попробуйте переформулировать её через Union-Find. Заметьте, как код сокращается, производительность растёт, а уверенность — тоже.

Какую задачу на связность вы покорите с помощью DSU? Напишите в комментариях — продолжим квест!