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)$ findfind/unionunion.
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, findfind and unionunion 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

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 mergemerge 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))

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

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 LL 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

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)$	—	`findfind`/`unionunion`: $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	`LL` = 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.

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.

If this helped you, consider buying me a coffee ☕

Buy me a coffee