Union-Find, Trie, Fenwick Tree & Segment Tree

Decision Guide: Which Structure to Use?

flowchart TD
    Q1{Connectivity or group membership?}
    Q1 -->|Yes| UF[Union-Find]
    Q1 -->|No| Q2{Strings or prefix queries?}
    Q2 -->|Yes| TR[Trie]
    Q2 -->|No| Q3{Range queries on array?}
    Q3 -->|No| Q4[Array or HashMap is sufficient]
    Q3 -->|Yes| Q5{Point updates and range sum only?}
    Q5 -->|Yes| FW[Fenwick Tree - BIT]
    Q5 -->|No| Q6{Range updates or min/max/GCD?}
    Q6 -->|Yes| ST[Segment Tree]
    Q6 -->|No| SP[Sparse Table - static O-1 query]
                            

Union-Find (Disjoint Set Union)

Union-Find Implementation

class UnionFind:
    """
    Union-Find with path compression + union by rank.
    All operations: O(alpha(n)) amortized ~ O(1) practical.
    """
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.components = n  # number of distinct groups

    def find(self, x):
        """Path compression: flatten the tree on every find."""
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # path compression
        return self.parent[x]

    def union(self, x, y):
        """Union by rank: attach smaller tree under larger."""
        px, py = self.find(x), self.find(y)
        if px == py:
            return False  # already connected
        if self.rank[px] < self.rank[py]:
            px, py = py, px
        self.parent[py] = px
        if self.rank[px] == self.rank[py]:
            self.rank[px] += 1
        self.components -= 1
        return True  # merged two different groups

    def connected(self, x, y):
        return self.find(x) == self.find(y)


# Number of Islands (LeetCode 200) using Union-Find
def numIslands(grid):
    rows, cols = len(grid), len(grid[0])
    uf = UnionFind(rows * cols)
    water_count = 0

    for r in range(rows):
        for c in range(cols):
            if grid[r][c] == '0':
                water_count += 1
            else:
                for dr, dc in [(0,1),(1,0)]:
                    nr, nc = r+dr, c+dc
                    if 0 <= nr < rows and 0 <= nc < cols and grid[nr][nc] == '1':
                        uf.union(r*cols+c, nr*cols+nc)

    return uf.components - water_count


grid = [['1','1','0'],['0','1','0'],['0','0','1']]
print("Islands:", numIslands(grid))  # 2

Union-Find Applications

Trie (Prefix Tree)

Trie Implementation

class TrieNode:
    def __init__(self):
        self.children = {}   # char -> TrieNode
        self.is_end = False
        self.count = 0       # words that pass through this node


class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        """O(L) — insert word into trie."""
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
            node.count += 1
        node.is_end = True

    def search(self, word):
        """O(L) — True if word exists in trie."""
        node = self.root
        for ch in word:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return node.is_end

    def starts_with(self, prefix):
        """O(P) — True if any word starts with prefix."""
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return True

    def count_starts_with(self, prefix):
        """O(P) — count of words with given prefix."""
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return 0
            node = node.children[ch]
        return node.count


# Demo
t = Trie()
for w in ['apple', 'app', 'apply', 'apt', 'banana']:
    t.insert(w)

print(t.search('app'))           # True
print(t.search('ap'))            # False (not a full word)
print(t.starts_with('ap'))       # True
print(t.count_starts_with('app')) # 3 (app, apple, apply)
print(t.count_starts_with('ban')) # 1

Trie Applications

Fenwick Tree (Binary Indexed Tree)

Fenwick Tree Implementation

class FenwickTree:
    """
    Binary Indexed Tree (BIT) for prefix sums.
    1-indexed. O(log n) update and query.
    """
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (n + 1)  # 1-indexed; index 0 unused

    def update(self, i, delta):
        """Add delta to position i (1-indexed). O(log n)."""
        while i <= self.n:
            self.tree[i] += delta
            i += i & (-i)   # move to parent by adding lowest set bit

    def query(self, i):
        """Prefix sum from 1 to i (inclusive). O(log n)."""
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & (-i)   # move to previous range by removing lowest set bit
        return s

    def range_query(self, l, r):
        """Sum from l to r inclusive. O(log n)."""
        return self.query(r) - self.query(l - 1)


# Count of Smaller Numbers After Self (LeetCode 315)
def countSmaller(nums):
    rank = {v: i+1 for i, v in enumerate(sorted(set(nums)))}
    n = len(rank)
    bit = FenwickTree(n)
    result = []
    for num in reversed(nums):
        r = rank[num]
        result.append(bit.query(r - 1))  # count of elements < num inserted so far
        bit.update(r, 1)
    return result[::-1]


print(countSmaller([5, 2, 6, 1]))  # [2, 1, 1, 0]

Segment Tree

Segment Tree Implementation

class SegmentTree:
    """
    Segment Tree for range sum queries with point updates.
    Build: O(n), Update: O(log n), Query: O(log n).
    """
    def __init__(self, nums):
        self.n = len(nums)
        self.tree = [0] * (4 * self.n)
        self._build(nums, 0, 0, self.n - 1)

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

    def update(self, idx, val, node=0, start=0, end=None):
        """Update index idx to value val. O(log n)."""
        if end is None: end = self.n - 1
        if start == end:
            self.tree[node] = val
        else:
            mid = (start + end) // 2
            if idx <= mid:
                self.update(idx, val, 2*node+1, start, mid)
            else:
                self.update(idx, val, 2*node+2, mid+1, end)
            self.tree[node] = self.tree[2*node+1] + self.tree[2*node+2]

    def query(self, l, r, node=0, start=0, end=None):
        """Range sum [l, r]. O(log n)."""
        if end is None: end = self.n - 1
        if r < start or end < l:
            return 0  # out of range
        if l <= start and end <= r:
            return self.tree[node]  # fully within range
        mid = (start + end) // 2
        left = self.query(l, r, 2*node+1, start, mid)
        right = self.query(l, r, 2*node+2, mid+1, end)
        return left + right


# Demo: Range Sum Query - Mutable (LeetCode 307)
st = SegmentTree([1, 3, 5, 7, 9, 11])
print(st.query(1, 3))   # sum [3,5,7] = 15
st.update(1, 10)        # set index 1 to 10
print(st.query(1, 3))   # sum [10,5,7] = 22

Comprehensive Comparison

Interview & Practice