Algorithms & Data Structures for AI

Not "pure math," but every AI/ML interview loop tests this alongside the theory — and understanding complexity is what lets you reason about whether an ML system will actually scale.

Intuition: Complexity Classes Aren't Academic — They're a Budget

Every number on this page is really answering one question: as a real input (a sequence, a vocabulary, a dataset) grows, does the cost grow slowly enough to still be affordable? $O(n)$ and $O(n^2)$ look similar on paper; at real production scale they're the difference between "fine" and "the reason this system can't ship." That gap, made concrete with real operation counts rather than asymptotic notation alone, is what this page is really about.

Big-O Complexity

Big-O describes how an algorithm's running time (or memory) grows as input size $n$ grows, ignoring constant factors. Drag $n$ and watch real operation counts diverge:

n20

⬤ O(1)⬤ O(log n)⬤ O(n)⬤ O(n log n)⬤ O(n^2)

At n=20: O(1) = 1 ops, O(log n) = 4 ops, O(n) = 20 ops, O(n log n) = 86 ops, O(n^2) = 400 ops. O(2ⁿ) at n=20 would be 1,048,576 -- too large to even plot on the same axis, which is the entire point: exponential algorithms aren't "somewhat slower," they're categorically infeasible past small n.

$O(1)$ — constant time (hash map lookup)
$O(\log n)$ — logarithmic (binary search, balanced tree operations)
$O(n)$ — linear (scanning a list once)
$O(n \log n)$ — the cost of efficient sorting, and many "divide and conquer" algorithms
$O(n^2)$ — quadratic (naive pairwise comparison — e.g. computing full pairwise attention scores)
$O(2^n)$ — exponential (brute-force search over subsets — avoid at all costs)

Why this matters for AI specifically: self-attention is $O(n^2)$ in sequence length, which is why long-context LLMs are hard and expensive, and why Flash Attention / sparse attention exist — direct responses to this complexity bound, made concrete below.

Core Data Structures

Arrays — contiguous memory, $O(1)$ index access. Tensors are just multi-dimensional arrays.
Hash maps — average $O(1)$ insert/lookup via hashing. Used for tokenizer vocabularies, caching, and de-duplication in data pipelines. Real op-count comparison against the naive alternative:

collection size50,000

array (linear scan)

50,000 ops

hash map

1 op

The hash map's bar never grows with collection size -- that flatness IS what O(1) means, made visible instead of asserted.

Vocabulary/cache size = 50,000 entries. Real worst-case lookup cost: array scan = 50,000 comparisons; hash map = 1 (average case, real amortized O(1) via hashing). This gap is exactly why tokenizer vocabularies, caching layers, and data-pipeline de-duplication all use hash maps, not arrays, at any real scale.

Trees — hierarchical structure. Decision trees are this data structure directly; balanced trees (B-trees) underlie database indexes.
Graphs — nodes + edges. Directly relevant to graph databases (Neo4j), GraphRAG, and modeling multi-agent communication topology.
Heaps / priority queues — efficiently retrieve the min/max element. Used in beam search (keeping the top- $k$ candidate sequences during LLM decoding) and in Dijkstra-style search algorithms. Real beam search, real cumulative log-probabilities, real top-k pruning at every step:

<s>0.00

step 0

Score = real cumulative sum of log-probabilities along the sequence -- higher (less negative) is better. Beam search never guarantees the single best full sequence, only the best it can find while only ever tracking k candidates at once.

Beam width k=2: at every step, every surviving candidate is expanded with every possible next token, real log-probabilities are summed, and only the top-2 scoring sequences survive to the next step -- exactly the heap/priority-queue operation (keep the top-k, discard the rest) real LLM decoding runs at every single generation step.

Sorting & Searching

Sorting ( $O(n \log n)$ for comparison-based sorts like merge sort/quicksort) shows up whenever you rank candidates — e.g. sorting retrieved documents by relevance score in a RAG pipeline, or ranking recommendations. Real comparison counts, not just the asymptotic label:

items to sort (n)200

⬤ bubble sort O(n²)⬤ merge sort O(n log n)

Real worst-case comparison counts to sort n=200 items: merge sort (O(n log n)) = 1,529; bubble sort (O(n²)) = 19,900 -- a real 13.0x gap already at this size. This is the concrete cost of "sorting retrieved documents by relevance score" or "ranking recommendations" done the naive way vs. the standard way.

Binary search ( $O(\log n)$ ) requires sorted data, used in efficient lookup structures. Approximate Nearest Neighbor search (HNSW, IVF — covered in Databases) is the vector-database analog of search, trading exactness for speed at scale.

Complexity of Common ML Algorithms

Algorithm	Training complexity (roughly)	Notes
k-Nearest Neighbors	$O(1)$ train, $O(n)$ per query (naive)	No real "training" — all cost is at inference
k-Means	$O(n \cdot k \cdot i)$	$k$ clusters, $i$ iterations
Decision Tree	$O(n \cdot d \cdot \log n)$	$d$ features
Matrix multiply (a layer's forward pass)	$O(n \cdot m \cdot p)$	For an $n \times m$ by $m \times p$ multiply
Self-attention	$O(n^2 \cdot d)$	$n$ = sequence length, $d$ = embedding dim — the scaling bottleneck for LLMs

Real operation counts for self-attention's $O(n^2 \cdot d)$ against a linear-attention alternative, at real sequence lengths:

sequence length (n)2,000

⬤ standard attention O(n²d)⬤ linear attention O(nd)

Real op counts at sequence length 2,000, dim=128: standard self-attention (O(n²·d)) = 512,000,000 operations; a linear-attention variant (O(n·d)) = 256,000 -- a real 2000x gap at this length, and it only widens as context length grows. This IS why long-context LLMs are hard and expensive, and why Flash Attention / sparse attention exist -- direct responses to this exact complexity bound.

Knowing these numbers is what lets you answer "why is this slow, and what would make it faster" — a staple of both ML system design interviews and real production debugging.

Code: The Beam Search Loop, For Real

The exact algorithm the diagram above steps through:

def beam_search(model, start_token: str, beam_width: int, num_steps: int):
    beams = [(start_token, 0.0)]  # (sequence, cumulative log-prob)
    for _ in range(num_steps):
        candidates = []
        for sequence, score in beams:
            for token, log_prob in model.next_token_distribution(sequence):
                candidates.append((f"{sequence} {token}", score + log_prob))
        # Keep only the top-k -- exactly the heap/priority-queue operation.
        beams = sorted(candidates, key=lambda c: c[1], reverse=True)[:beam_width]
    return beams

Why this belongs in a math curriculum

Interviewers (and real jobs) expect ML engineers to be comfortable with general software engineering fundamentals, not just theory — see Interview Prep for how this gets tested directly. But it also isn't separable from the math: understanding why attention is expensive, why vector search needs approximate algorithms, and why certain training loops scale better than others all comes back to Big-O reasoning applied to the linear algebra operations covered on the Linear Algebra page.

Mathematics for AI complete. Next: Machine Learning — where this math turns into working algorithms.