Data Structure Templates

This page is a cheat sheet, not a tutorial — four self-contained, runnable templates you can copy straight into a contest editor or an interview whiteboard. Each one is battle-tested, uses only the standard library, and ends with a tiny demo so you can sanity-check it the moment you paste it in.

What you’ll learn

Union-Find (DSU) with path compression + union by rank — near $O(1)$ find/union.
Segment Tree with point update + range query — $O(\log n)$ for any associative merge (sum, min, max).
Fenwick Tree (BIT) — the same range-sum-and-point-update problem in a third of the code.
Trie — $O(L)$ insert/search/prefix-check, independent of how many words are stored.

Union-Find (DSU)

Use it whenever a problem merges things into groups over time and asks “same group?” — connected components, Kruskal’s MST, cycle detection on an undirected graph. With path compression and union by rank, find and union both run in $O(\alpha(n))$ amortized — effectively $O(1)$ .

dsu_template.py

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.count = n   # number of distinct components
 
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])   # path compression
        return self.parent[x]
 
    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.rank[ra] < self.rank[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        if self.rank[ra] == self.rank[rb]:
            self.rank[ra] += 1
        self.count -= 1
        return True
 
 
dsu = DSU(6)
for a, b in [(0, 1), (1, 2), (3, 4)]:
    dsu.union(a, b)
 
print("components:", dsu.count)                 # {0,1,2}, {3,4}, {5} -> 3
print("0 and 2 connected?", dsu.find(0) == dsu.find(2))   # True
print("0 and 3 connected?", dsu.find(0) == dsu.find(3))   # False

Segment Tree (point update, range query)

Use it when a range query (sum/min/max/gcd/…) and point updates are both needed, repeatedly, in any order. Iterative and array-based — swap merge for any associative function. $O(n)$ build, $O(\log n)$ per update or query.

segment_tree_template.py

class SegmentTree:
    def __init__(self, data, merge=lambda a, b: a + b, identity=0):
        self.n = len(data)
        self.merge = merge
        self.identity = identity
        self.tree = [identity] * (2 * self.n)
        for i, v in enumerate(data):
            self.tree[self.n + i] = v
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = merge(self.tree[2 * i], self.tree[2 * i + 1])
 
    def update(self, pos, value):
        i = pos + self.n
        self.tree[i] = value
        while i > 1:
            i //= 2
            self.tree[i] = self.merge(self.tree[2 * i], self.tree[2 * i + 1])
 
    def query(self, left, right):
        """Half-open range [left, right)."""
        result = self.identity
        left += self.n
        right += self.n
        while left < right:
            if left & 1:
                result = self.merge(result, self.tree[left])
                left += 1
            if right & 1:
                right -= 1
                result = self.merge(result, self.tree[right])
            left //= 2
            right //= 2
        return result
 
 
data = [2, 4, 5, 7, 8, 9]
st = SegmentTree(data)                 # default: range sum
print("sum[1:4):", st.query(1, 4))     # 4 + 5 + 7 = 16
 
st.update(2, 10)                       # data[2] becomes 10
print("sum[1:4) after update:", st.query(1, 4))   # 4 + 10 + 7 = 21
 
min_st = SegmentTree(data, merge=min, identity=float("inf"))
print("min[0:6):", min_st.query(0, 6))

Fenwick Tree (Binary Indexed Tree)

Use it for the exact same range-sum-and-point-update problem, when the operation is invertible (sum, xor) — roughly a third of the code, no recursion. $O(\log n)$ per update/query, 1-indexed.

fenwick_tree_template.py

class FenwickTree:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (n + 1)
 
    def update(self, i, delta):
        while i <= self.n:
            self.tree[i] += delta
            i += i & (-i)
 
    def prefix_sum(self, i):
        total = 0
        while i > 0:
            total += self.tree[i]
            i -= i & (-i)
        return total
 
    def range_sum(self, left, right):
        return self.prefix_sum(right) - self.prefix_sum(left - 1)
 
    @classmethod
    def from_list(cls, values):
        bit = cls(len(values))
        for i, v in enumerate(values, start=1):
            bit.update(i, v)
        return bit
 
 
values = [3, 2, -1, 6, 5, 4, -3, 3]   # 1-indexed: position 1..8
bit = FenwickTree.from_list(values)
 
print("sum of first 5:", bit.prefix_sum(5))       # 3+2-1+6+5 = 15
print("sum of positions 3..6:", bit.range_sum(3, 6))   # -1+6+5+4 = 14
 
bit.update(3, 10)                                  # position 3: -1 -> 9
print("sum of 3..6 after update:", bit.range_sum(3, 6))   # 9+6+5+4 = 24

Trie (Prefix Tree)

Use it for prefix-aware string queries — autocomplete, “does any word start with X”, word search on a board. $O(L)$ per insert/search/prefix check, where L is the string length — independent of how many words are stored.

trie_template.py

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False
 
 
class Trie:
    def __init__(self):
        self.root = TrieNode()
 
    def insert(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True
 
    def _walk(self, prefix):
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node
 
    def search(self, word):
        node = self._walk(word)
        return node is not None and node.is_end
 
    def starts_with(self, prefix):
        return self._walk(prefix) is not None
 
 
trie = Trie()
for word in ["cat", "car", "dog"]:
    trie.insert(word)
 
print("search 'cat':", trie.search("cat"))            # True
print("search 'ca':", trie.search("ca"))               # False -- only a prefix
print("starts_with 'ca':", trie.starts_with("ca"))     # True

Complexity at a glance

Structure	Build	Point update	Range/prefix query	Notes
Union-Find	$O(n)$	—	`find`/`union`: $O(\alpha(n))$ amortized	Groups, not ranges
Segment Tree	$O(n)$	$O(\log n)$	$O(\log n)$	Any associative merge
Fenwick Tree	$O(n \log n)$	$O(\log n)$	$O(\log n)$	Invertible ops only (sum, xor)
Trie	$O(n \cdot L)$ total	$O(L)$ insert	$O(L)$ search/prefix	`L` = word length

Practice

Drill 1 — DSU union by rank. Complete the swap that keeps the taller tree on top.

Drill 2 — Fenwick tree’s lowbit walk. Complete the query step that walks toward index 0, one disjoint range at a time.

Union-Find — count connected components

Pitfalls

Sizing a recursive segment tree 2n. The iterative bottom-up layout fits in 2n; the recursive one needs 4n, because a non-power-of-two n pushes node indices past 2n. It is an out-of-range error that only appears for certain sizes.
Decrementing the DSU component count on a failed union. union returning False means the two were already connected — the count must not change. Getting this wrong under-reports components by the number of redundant edges.
A 0-indexed Fenwick tree. lowbit(0) == 0, so an update at index 0 loops forever and a query at 0 returns immediately. tree[0] is unused and every public call is 1-indexed.
Using a Fenwick tree for range min/max. It needs an invertible operation, because a range query is prefix(r) - prefix(l-1). Min has no inverse — that is a segment tree.
Forgetting to push lazy tags down before descending. Children then return values from before the pending range update: a wrong answer with no crash.
A trie built with a plain dict and no end-of-word marker. Without it, "car" matches inside "cart" — search and startsWith become the same function.
Path compression written recursively on a deep chain. find recursing $10^5$ deep hits CPython’s frame limit; the iterative path-halving form in these templates avoids it.
Reaching for a segment tree when the array never changes. Static range min is a sparse table at $O(1)$ per query; static range sum is a prefix array.

Interview follow-ups

They ask	What they’re checking	The answer
“Union-find or DFS for counting components?”	Judgement	Identical for a static graph — DFS is $O(V+E)$ and simpler. Union-find wins when edges arrive over time and connectivity is queried between arrivals, at $O(\alpha)$ per operation instead of a full re-traversal
“What is union-find’s real complexity?”	Precision	Effectively $O(1)$ ; formally $O(\alpha(n))$ amortised with both path compression and union by rank. With only one of them it degrades to $O(\log n)$
“Fenwick or segment tree?”	Choosing correctly	Fenwick for invertible operations with point updates — a third of the code and better constants. Segment tree for min/max/gcd, or when you need range updates with lazy propagation
“What does `tree[i]` hold in a Fenwick tree?”	Whether you understand it or memorised loops	The sum of the `lowbit(i)` elements ending at `i`. Update walks up the ranges containing `i`; query walks down the disjoint ranges tiling the prefix
“How much memory does a trie use?”	Practical cost	$O(\text{total characters})$ nodes, each with a children map — heavy in Python. For a fixed lowercase alphabet a 26-slot list per node is faster but wastes space; a `dict` is the usual compromise
“Make the segment tree support range assignment”	Lazy propagation	Add a `lazy[]` array, mark the covering node, and push down before descending. Assignment and addition need different push-down rules — assignment overwrites a child’s tag, addition accumulates
“Which of these would you write from memory in an interview?”	Honesty	Union-find, unquestionably — it is ten lines and appears constantly. A segment tree is worth deriving rather than recalling, and a Fenwick tree is short enough to rebuild from the `lowbit` idea
“The keys are strings, not integers”	Adapting	A trie indexes by character directly. For a Fenwick or segment tree, coordinate-compress the keys to ranks first — sort the distinct values and index by position

LeetCode problem set

The problems these four structures exist for. If you cannot tell at a glance which of union-find, segment tree, Fenwick tree or trie a row wants, that is the gap to close.

23 problems

0 easy18 medium5 hard

Work down the ladder. Tick each problem off as you go — progress is saved in this browser, and the Export button in the filter bar writes it to a file you can keep.

Unsolved onlyImport

208Implement Trie (Prefix Tree)mediumNeetCode 150Blind 75LeetCode Top Interview 150googleamazonbloomberg
128Longest Consecutive SequencemediumNeetCode 150Blind 75LeetCode Top Interview 150
211Design Add and Search Words Data StructuremediumNeetCode 150Blind 75LeetCode Top Interview 150
261Graph Valid TreepremiummediumNeetCode 150Blind 75
307Range Sum Query - Mutablemedium
308Range Sum Query 2D - Mutablepremiummedium
399Evaluate DivisionmediumLeetCode Top Interview 150
421Maximum XOR of Two Numbers in an Arraymedium
547Number of Provincesmedium
648Replace Wordsmedium
673Number of Longest Increasing Subsequencemedium
677Map Sum Pairsmedium
684Redundant ConnectionmediumNeetCode 150
720Longest Word in Dictionarymedium
721Accounts Mergemedium
990Satisfiability of Equality Equationsmedium
1268Search Suggestions Systemmedium
1319Number of Operations to Make Network Connectedmedium
212Word Search IIhardNeetCode 150Blind 75LeetCode Top Interview 150googleamazonbytedance
315Count of Smaller Numbers After Selfhard
699Falling Squareshard
732My Calendar IIIhard
1649Create Sorted Array through Instructionshard

Recall card

Union-Find — parent + rank, path-halving find, union returns False when already joined. Keep a components counter and decrement only on a successful union. $O(\alpha(n))$ .
Segment tree — flat array, leaves at n..2n-1, node i = merge(2i, 2i+1). Any associative merge. 2n iterative, 4n recursive.
Fenwick tree — tree[i] = sum of the lowbit(i) elements ending at i. Update i += i & -i, query i -= i & -i. 1-indexed always. Invertible operations only.
Trie — a dict per node plus an end-of-word flag; without the flag search and startsWith collapse into one thing.
Choosing — static + min/max → sparse table; mutable + invertible → Fenwick; mutable + min/max or range updates → segment tree; connectivity over arriving edges → union-find.
All four are worth typing once from memory. Union-find is the only one short enough that you should.

Recap

Union-Find: near- $O(1)$ group membership; reach for it on connected components and Kruskal’s MST.
Segment Tree: $O(\log n)$ range query + point update for any associative merge — the most general of the four.
Fenwick Tree: the same problem, a third of the code, when the operation is invertible.
Trie: $O(L)$ prefix-aware string queries, independent of dictionary size.

Next: Graph Algorithm Templates — copy-paste-ready BFS, DFS, Dijkstra, Prim, Kruskal, and Kahn’s topological sort.

pch.coffeeTagline

pch.coffeeCta

pch.feedbackHeading

pch.feedbackSubheading