Union-Find & Advanced Structures

Imagine you are managing a social network and need to answer: "Are Alice and Bob in the same friend group?" and "Merge Alice's group with Bob's group." The naive approach — maintaining explicit lists — costs O(n) per merge. Union-Find solves this with a forest of trees, one tree per group, where every node stores a pointer to its parent. The root of each tree is the canonical representative of its group.

The two operations are:

• find(x): walk up the parent pointers until you reach the root. Returns the group's representative.
• union(x, y): find the roots of x and y. If different, make one root the parent of the other — merging two trees into one.

Without optimizations, find is O(depth), and in the worst case (union a chain) the tree degenerates to a linked list with depth O(n). The magic is in the two optimizations that together make this near-constant.

# Bare-bones Union-Find (no optimizations yet)
class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))  # parent[i] = i means i is its own root

    def find(self, x):
        while self.parent[x] != x:
            x = self.parent[x]
        return x

    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx != ry:
            self.parent[rx] = ry  # make ry the new root

The first optimization is path compression. Every time we call find(x), we walk from x up to the root. After we find the root, why leave all those intermediate nodes pointing to their old parents? They could all point directly to the root — the next call to find on any of them would then be O(1).

Path compression does exactly this: on the way back from the recursive call (or in a two-pass iterative version), we set parent[x] = root for every node x along the path. This is sometimes called full path compression. A simpler variant, path splitting, makes each node on the path point to its grandparent — only one write per node — and achieves the same asymptotic complexity.

Path compression alone (without union by rank) gives amortized O(log n) per operation. It's already a huge win, but the real magic comes when we combine it with union by rank.

def find(self, x):
    if self.parent[x] != x:
        self.parent[x] = self.find(self.parent[x])  # path compression
    return self.parent[x]

# Iterative path splitting (often preferred in practice — no recursion stack):
def find_split(self, x):
    while self.parent[x] != x:
        self.parent[x] = self.parent[self.parent[x]]  # point to grandparent
        x = self.parent[x]
    return x

The second optimization addresses which root becomes the parent when we merge two trees. Without care, always attaching a tall tree under a short one explodes the depth. Union by rank maintains an upper bound on tree height called the rank. When merging, the root with lower rank is attached under the root with higher rank. If ranks are equal, attach either way and increment the surviving root's rank by 1.

Crucially, rank only represents the height when there is no path compression. Once path compression fires, trees get flatter but we don't bother updating ranks (it would be expensive). The rank becomes an upper bound on depth, not the exact height.

Union by size is an alternative: track the number of nodes in each subtree, attach the smaller tree under the larger. This gives the same O(log n) depth guarantee and is sometimes simpler to reason about (actual subtree size is a useful extra piece of information).

With either of these union strategies alone (no path compression), the depth is bounded by O(log n), giving O(log n) per operation. The two optimizations together produce something far better.

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x):  # with path compression
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return False  # already in same set
        # union by rank: attach smaller-rank tree under larger
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1
        return True

Combined, path compression + union by rank gives amortized O(α(n)) per operation, where α is the inverse Ackermann function. This is the most famous nearly-constant complexity in computer science, and understanding it deserves a moment.

The Ackermann function A(m, n) grows obscenely fast — faster than any tower of exponentials, faster than any primitive recursive function. It is the quintessential function that escapes all "reasonable" growth patterns. Its inverse, α(n), tells you: for what m does A(m, m) first exceed n? Because A grows so fast, α grows incredibly slowly. For n = 2^65536 (a number with roughly 10^19000 digits), α(n) is still just 4. For all practical purposes, α(n) ≤ 5 for any n you will ever encounter.

This result, proved by Tarjan in 1975, shows that Union-Find with both optimizations is not just O(log n) but essentially O(1) per operation in any real application. The proof is non-trivial (it uses a potential function argument and a clever accounting of path-compression work across operations), but the conclusion is profound: a data structure that seems like it should be O(log n) is, with two simple tricks, indistinguishable from O(1).

For a sequence of m operations on n elements, the total cost is O(m · α(n)) — essentially O(m).

Suppose you have an array A of n numbers and need to repeatedly answer: "What is the sum of A[l..r]?" and "Set A[i] = v." A prefix-sum array answers sum queries in O(1) but updates in O(n). A flat array updates in O(1) but queries in O(n). A segment tree achieves both in O(log n).

The segment tree is a complete binary tree where:

• Each leaf holds one array element.
• Each internal node holds the aggregate (sum, min, max, GCD — any associative operator) of all elements in its subtree's range.
• The root covers the entire array; the two children split the range in half.

Building the tree takes O(n). Query: to sum A[l..r], recurse down the tree, collecting nodes whose ranges are fully inside [l, r] and ignoring those fully outside. At most O(log n) nodes are visited. Update: change a leaf, then update all O(log n) ancestors.

The key insight is that each node's range is determined purely by its position in the tree — you do not need to store ranges explicitly. With 1-indexed nodes, node i's children are 2i and 2i+1, and the tree fits in an array of size 4n.

# Segment tree for range sum queries, point updates
class SegTree:
    def __init__(self, arr):
        self.n = len(arr)
        self.tree = [0] * (4 * self.n)
        self._build(arr, 1, 0, self.n - 1)

    def _build(self, arr, node, start, end):
        if start == end:
            self.tree[node] = arr[start]
        else:
            mid = (start + end) // 2
            self._build(arr, 2*node,   start, mid)
            self._build(arr, 2*node+1, mid+1, end)
            self.tree[node] = self.tree[2*node] + self.tree[2*node+1]

    def update(self, node, start, end, idx, val):
        if start == end:
            self.tree[node] = val
        else:
            mid = (start + end) // 2
            if idx <= mid:
                self.update(2*node,   start, mid,   idx, val)
            else:
                self.update(2*node+1, mid+1, end,   idx, val)
            self.tree[node] = self.tree[2*node] + self.tree[2*node+1]

    def query(self, node, start, end, l, r):
        if r < start or end < l:    # completely outside
            return 0
        if l <= start and end <= r: # completely inside
            return self.tree[node]
        mid = (start + end) // 2
        return (self.query(2*node,   start, mid,   l, r) +
                self.query(2*node+1, mid+1, end,   l, r))

Plain segment trees handle point updates. But what if you want to add 5 to every element in A[3..17]? Naive point-update for each element costs O(n log n). Lazy propagation reduces this to O(log n) by deferring work.

The idea: when we update a range, instead of immediately updating all affected nodes, we store a pending (lazy) tag at the highest nodes whose ranges are completely covered. When we later need to visit a node's children (for a query or a deeper update), we first push down the pending tag to the children, then clear it from the current node.

This ensures that: (a) the node's own aggregate value is always up-to-date before we use it, and (b) pending work is only materialized when it becomes necessary — O(log n) nodes per operation.

The lazy mechanism is customizable. For range-add + range-sum queries, the tag is an additive offset. For range-assign + range-min queries, the tag is the assigned value. Composing two different update types (add and assign together) requires careful tag composition and is a classic exercise.

Lazy propagation is one of the most powerful techniques in competitive programming — it turns segment trees into a Swiss Army knife for range query problems.

# Segment tree with lazy propagation: range add, range sum query
class LazySegTree:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (4 * n)
        self.lazy = [0] * (4 * n)  # pending additions

    def _push_down(self, node, start, end):
        if self.lazy[node] != 0:
            mid = (start + end) // 2
            # apply lazy to children
            self.tree[2*node]   += self.lazy[node] * (mid - start + 1)
            self.tree[2*node+1] += self.lazy[node] * (end - mid)
            self.lazy[2*node]   += self.lazy[node]
            self.lazy[2*node+1] += self.lazy[node]
            self.lazy[node] = 0

    def update(self, node, start, end, l, r, val):
        if r < start or end < l: return
        if l <= start and end <= r:
            self.tree[node] += val * (end - start + 1)
            self.lazy[node] += val
            return
        self._push_down(node, start, end)
        mid = (start + end) // 2
        self.update(2*node,   start, mid,   l, r, val)
        self.update(2*node+1, mid+1, end,   l, r, val)
        self.tree[node] = self.tree[2*node] + self.tree[2*node+1]

    def query(self, node, start, end, l, r):
        if r < start or end < l: return 0
        if l <= start and end <= r: return self.tree[node]
        self._push_down(node, start, end)
        mid = (start + end) // 2
        return (self.query(2*node,   start, mid,   l, r) +
                self.query(2*node+1, mid+1, end,   l, r))

Peter Fenwick's 1994 paper introduced a data structure so compact it seems almost magical: a flat array that answers prefix sum queries (sum of A[1..i]) and handles point updates in O(log n), using only n extra integers and perhaps 10 lines of code.

The key insight exploits the binary representation of indices. Every positive integer i can be written as a sum of powers of 2. The least significant set bit of i, written lowbit(i) = i & (-i), tells you how large a range the BIT node at position i is responsible for: node i stores the sum of elements from i - lowbit(i) + 1 to i.

Prefix sum query (sum of A[1..i]): start at i, add BIT[i], then jump to i - lowbit(i), repeat until i = 0. You are stripping off the lowest set bit each time — this terminates in O(log n) steps.

Point update (add delta to A[i]): start at i, add delta to BIT[i], then jump to i + lowbit(i), repeat until i > n. You are propagating up to all the nodes responsible for a range containing i — O(log n) steps.

The elegance is in the symmetry: query walks down by removing bits, update walks up by adding bits. The two operations are perfect inverses.

For range-sum queries on [l, r], use prefix(r) - prefix(l-1). The BIT is 1-indexed by convention.

class FenwickTree:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)  # 1-indexed

    def update(self, i, delta):
        """Add delta to position i (1-indexed)."""
        while i <= self.n:
            self.bit[i] += delta
            i += i & (-i)          # i + lowbit(i)

    def prefix(self, i):
        """Sum of A[1..i]."""
        s = 0
        while i > 0:
            s += self.bit[i]
            i -= i & (-i)          # i - lowbit(i)
        return s

    def range_sum(self, l, r):
        return self.prefix(r) - self.prefix(l - 1)

# Build from an existing array in O(n log n):
def build(arr):
    ft = FenwickTree(len(arr))
    for i, v in enumerate(arr, 1):
        ft.update(i, v)
    return ft

# O(n) build trick: direct parent propagation
def build_fast(arr):
    n = len(arr)
    bit = [0] + arr[:]          # 1-indexed copy
    for i in range(1, n + 1):
        j = i + (i & -i)
        if j <= n:
            bit[j] += bit[i]
    return bit

A sorted linked list supports O(n) search — you have to scan from the front. A balanced BST (AVL, Red-Black) supports O(log n) but requires complex rotations to maintain balance. Skip lists, invented by William Pugh in 1990, achieve O(log n) expected time for search, insert, and delete using only randomization and a tower of linked lists — no rotations, no rebalancing.

The structure: imagine stacking multiple linked lists on top of each other, each a sparser version of the one below. The bottom level (level 0) contains all elements in sorted order. Each element is promoted to level 1 with probability 1/2, to level 2 with probability 1/4, and so on. The expected number of levels is O(log n).

Search for key k: start at the top-left. Move right as long as the next key is ≤ k. When you can't move right, drop down one level. Repeat until you hit level 0. You've found k (or established it's absent). Think of the upper levels as express lanes on a highway: you zoom past large chunks at high levels, then refine at lower levels.

Insert: find the insertion point (as in search, tracking where you drop down at each level). Roll a coin repeatedly to determine the new node's height. Link it into each level it participates in.

Expected O(log n) analysis: At level k, roughly n/2^k elements are present. A search at level k traverses O(1) elements on average before dropping down (geometric distribution). Summing over O(log n) levels gives O(log n) expected total steps.

The worst case is O(n) — all n elements could be promoted to the top level — but this probability is exponentially small and negligible in practice.

import random

class SkipNode:
    def __init__(self, key, val, height):
        self.key = key
        self.val = val
        self.forward = [None] * height  # pointers at each level

class SkipList:
    MAX_LEVEL = 16
    P = 0.5

    def __init__(self):
        self.head = SkipNode(float('-inf'), None, self.MAX_LEVEL)
        self.level = 1  # current max level in use

    def _random_level(self):
        lvl = 1
        while random.random() < self.P and lvl < self.MAX_LEVEL:
            lvl += 1
        return lvl

    def search(self, key):
        cur = self.head
        for i in range(self.level - 1, -1, -1):   # top to bottom
            while cur.forward[i] and cur.forward[i].key < key:
                cur = cur.forward[i]
        cur = cur.forward[0]
        return cur.val if cur and cur.key == key else None

    def insert(self, key, val):
        update = [self.head] * self.MAX_LEVEL
        cur = self.head
        for i in range(self.level - 1, -1, -1):
            while cur.forward[i] and cur.forward[i].key < key:
                cur = cur.forward[i]
            update[i] = cur
        cur = cur.forward[0]
        if cur and cur.key == key:
            cur.val = val  # update existing
        else:
            new_level = self._random_level()
            if new_level > self.level:
                for i in range(self.level, new_level):
                    update[i] = self.head
                self.level = new_level
            new_node = SkipNode(key, val, new_level)
            for i in range(new_level):
                new_node.forward[i] = update[i].forward[i]
                update[i].forward[i] = new_node

Knowing these structures deeply means knowing when to reach for each:

Union-Find shines whenever you need dynamic connectivity: Kruskal's MST, offline LCA queries, detecting cycles in undirected graphs, percolation problems, network connectivity, and clustering. Its amortized near-O(1) per operation is unbeatable. It does not support disconnection (splitting a set into two is not supported efficiently).

Segment Trees are the right tool when you need O(log n) range queries and O(log n) point (or range, with lazy) updates over an array, and the operation is associative (sum, min, max, GCD, XOR, matrix multiply...). They support arbitrary aggregation and are extremely versatile. Use them when the constant factor of a Fenwick tree isn't worth the restricted API.

Fenwick Trees are ideal when you specifically need prefix sums (or any group operation — prefix XOR, prefix max if you only query prefix and not arbitrary ranges) with point updates. They are 2x more memory-efficient than segment trees, have smaller constants, and fit in ~10 lines. For competitive programming, if the problem fits, always try a Fenwick tree first.

Skip Lists are for ordered dictionaries (sorted maps/sets) with O(log n) insert/delete/search, especially when you need easy range traversal or the overhead of tree rotations is prohibitive. Redis's sorted sets use skip lists. They are also easier to make lock-free in concurrent settings than BSTs, which is why they appear in concurrent data structure literature.

A common trick: offline processing + DSU (process queries in reverse order, treating deletions as insertions). Segment trees on coordinates (coordinate compression before building the tree) handle sparse, large-coordinate problems. 2D Fenwick trees extend the bit trick to grids for 2D prefix sums.