Algorithms for ML Engineers

An internal engineering assistant receives one question:

Why did the auth tests fail?

Its permission filter returns four authorized text chunks (small passages) in a local fixture. Only one mentions the expired test fixture. Finding it still requires reading all four chunks. What happens when the same search has to inspect 50,000?

An algorithm is a recipe for completing a task; complexity describes how its work or memory grows as its input grows. We'll keep the four-row search visible while counting the work that a larger corpus would repeat.

What grows when the corpus grows

Before choosing a faster lookup, pin down what can grow. This lesson builds on the SQL permission filter: predicates, optionally reinforced by row-level security (per-row access rules), have already returned only auth project chunks permitted for this request. Four symbols name the inputs that later code will grow with:

A permission filter can shrink n, and a wider embedding increases d. Returning two citations controls k, but it doesn't skip scoring every remaining candidate. That last distinction will matter when we replace a full sort with a heap.

Count a complete scan

Start with an inspectable retriever. It scores a chunk by counting question words that appear in the chunk text. The four rows below stay the running corpus for the rest of the chapter. Before running the code, predict one detail: the question matches one row, but how many rows will the function inspect?

chunks_checked isn't elapsed time. It's a count of the dominant repeated action: inspecting one candidate chunk. Four authorized chunks produce four inspections; forty thousand produce forty thousand.

Real text work also depends on chunk length. A configured chunk-size limit lets us use candidate count as a useful first model. Without that bound, words read per chunk become another growing input. A counter exposes repeated work; it doesn't measure how quickly a machine performs it.

Find that repeated operation before reaching for an index.

Big-O names the growth shape

Big-O gives an asymptotic upper bound: how a cost can grow as an input gets large, ignoring fixed multipliers and small-input details. The label alone doesn't say whether you're describing a worst case, an average case, or an amortized sequence of operations, so name both the operation and the kind of bound.^{[1]Reference 1big-O notationhttps://xlinux.nist.gov/dads/HTML/bigOnotation.html}

Under the bounded-chunk simplification, this scan visits every authorized chunk. Its candidate-inspection work is tight: Θ(n), which also means O(n). Read Θ as "theta": it says the count grows linearly, not merely that linear growth is an upper limit.

Watch the count while repeating the same four rows. Repetition here changes workload size, not the meaning or quality of the documents:

Doubling n doubles the repeated checks. Building a set of question terms still costs something, but that overhead doesn't grow once per candidate. Big-O points to the piece that becomes dangerous at larger scale.

Don't call the whole assistant "O(n)." Attach the claim to an operation: the candidate-inspection part of this bounded-length scan is Θ(n). The SQL filter, result ordering, network calls, and answer generation have their own costs.

There is another operation hiding in the function: it sorts its r matching chunks before returning them. Here r is the number of matches, and r <= n. The scan is Θ(n), while that ordering step is O(r log r). If every chunk matches, the whole function can be O(n log n). Naming each piece keeps a fast scan from hiding an expensive sort.

Choose a bound that says what you mean

Three words keep complexity claims honest. They answer a different question from the growth class itself:

Θ is a tight growth class: both an upper and a lower bound hold, up to fixed multipliers for sufficiently large inputs.^{[2]Reference 2Thetahttps://xlinux.nist.gov/dads/HTML/theta.html} The pair loop later performs exactly n * (n - 1) / 2 comparisons, so its loop count is Θ(n²); writing O(n²) alone states only the upper bound.

O is still useful when you promise only an upper bound, but it isn't automatically a worst-case label. Average-case analysis depends on an input distribution, expected-case analysis averages randomness or hash behavior, and amortized analysis accounts for operation history.

Pick a search that matches the data

Once the claim is scoped, choose a lookup that matches the data. An exact key in an unsorted list, an exact key in a sorted list, and semantic similarity are different problems. The Data Structures for AI lesson matched those questions to containers; here the same choices are about work. An approximate-nearest-neighbor (ANN) index changes the candidate path, so its recall needs measurement alongside latency. In the graph row, V counts nodes and E counts links:

For 50,000 sorted IDs, binary search needs at most 16 interval decisions because $2^{16} = 65{,}536$. Sorting those IDs on every request would erase the benefit, so sorted order or a hash index must be built and maintained outside the hot path.

Binary search also can't rank embeddings. Semantic similarity doesn't provide one scalar ordering where nearby vectors are guaranteed to sit next to each other. That is a different input and a different index problem.

A hash set is the right tool when you already have the exact ID. It can't answer "which stored prompts start like this one?" That is a prefix problem. The Data Structures lesson built a character trie for token, tool, and top. This chapter counts that work, then compresses the trie into a structure serving engines actually use.

Binary search halves a sorted range

The four authorized chunk IDs are already sorted alphabetically. At index 2 sits cli-flag, so searching for that ID is a lucky one-step hit. On a miss, the ordering tells us which half can still contain the target.

lo is the first possible index and hi is one past the last. The remaining range is therefore [lo, hi). Integer division // 2 chooses a midpoint; moving a boundary discards the part that can't contain the target.

The search finishes in one midpoint check here. An early-stopping linear search would find cli-flag on its third check; the earlier ranking scan would inspect all four. Binary search's worst-case guarantee comes from halving the remaining interval, so its step count stays O(log n) even when n is tens of thousands. A midpoint check may perform both == and <; we're counting visited midpoints, not individual comparison operators.

Exact counts still need care. This half-open implementation can require 17 midpoint checks for 65,536 items because the final one-element interval needs its own comparison. The list has to stay sorted, and a new ID inserted in the middle of a Python list still costs O(n).

Prefix match, tries, and radix trees

Suppose several requests begin with the same system prompt and tool schema. A lookup that rereads that shared prefix for every stored request pays for the same tokens again and again. A trie stores one symbol per edge: lookup of a prefix of length p walks p edges, then enumeration costs whatever the returned completions require. Work to reach the prefix doesn't grow with unrelated keys elsewhere.

Long unique suffixes waste nodes. The words token and tokenize share token; tokenize then adds a chain i -> z -> e that never branches. A radix tree, also called a Patricia trie or compressed prefix tree, collapses each non-branching run into one edge that holds a sequence. Lookup still follows exact symbols, but it compares a whole shared run at once.

That's not radix sort, and it isn't nearest-neighbor search. A radix tree answers a longest exact prefix. An embedding index answers approximate nearby vectors. The names sound close enough to mix up, but the lookup contract is different.

Three stored prompts share the system and tool tokens, then split. The compressed edge stores system + tools once. A new request that starts the same way follows that edge before checking its remaining tokens. Keep that picture in mind while the next program counts the work with token IDs.

Count the difference on toy token IDs. 1, 2 is the shared system-and-tools prefix. 3 is user A. 4, 11 is user B's whole remaining path.

The naive scan re-reads 1, 2 on every stored prompt. Three keys make 11 token comparisons. The radix walk compares the shared edge once, then the 3 branch, and stops: 3 explicit token comparisons for the same longest prefix of length 3. It also makes three dictionary lookups, including the failed lookup for 13; those aren't included in the comparison counter.

Indexing query[index + offset] avoids copying the remaining query at each edge. Writing query[index:] inside the loop would create a new list every time. Many short edges could then cause quadratic copying even while the token-comparison counter stayed linear.

Add a thousand other prompts that share only [1, 2] and the scan keeps paying per key. With expected constant-time child lookup, the tree's descent takes O(p + 1) work for p matched tokens. The extra 1 accounts for detecting a mismatch, including an immediate miss. Enumeration after the match is a separate cost, because returning more completions means doing more work.

For token sequences $x$ and $y$, the longest common prefix is the greatest number $p$ of initial tokens that agree. In this notation, $x_{1:j}$ means the first $j$ tokens; $j=0$ includes the empty prefix:

$$p=\max\{j:x_{1:j}=y_{1:j}\}.$$

Matching tokens can make previously computed model state reusable, but only when the model and other inputs that produced that state also match. The unmatched suffix still needs processing; a prefix hit doesn't make the whole request free.

Serving uses the same idea on token IDs. A prefix cache can reuse the prompt-processing work, called prefill, for a compatible cached prefix. SGLang's RadixAttention organizes token prefixes in a radix tree.^{[3]Reference 3SGLang: Efficient Execution of Structured Language Model Programshttps://arxiv.org/abs/2312.07104} vLLM instead hashes full blocks of cached model state. Its keys include preceding blocks and additional identity such as adapter or multimodal inputs, and block boundaries limit the reused length. Exact token agreement is necessary, but doesn't by itself establish safe reuse across models or cache scopes.^{[4]Reference 4Automatic Prefix Cachinghttps://docs.vllm.ai/en/stable/design/prefix_caching/}

You don't need to implement the cache here. Keep its algorithm test: if the question is longest exact prefix, count work in p. If the question is nearby meaning, you're back in the $n \times d$ scoring column.

Vector scoring adds another input

Exact prefix matching can't answer a nearby-meaning question. Semantic retrieval commonly represents the question and each chunk as embedding vectors, then computes a score. The NumPy and tensor shapes lesson already treated a vector as coordinates and a dot product as paired multiplies plus adds. Here we count those multiplies in plain Python so the new input stays visible.

Each exact dot-product score inspects d coordinates. When both vectors have length one, their dot product equals cosine similarity. The lab vectors below aren't normalized that way, so treat this as a raw dot product. Scoring all n allowed chunks performs work proportional to n * d.

Before running the code, predict the counter rather than the winner: four vectors with three coordinates each require how many coordinate multiplications, even when two query coordinates are zero?

The question vector is [1.0, 0.0, 0.0], so only the first coordinate of each chunk contributes to each sum. The other multiplies still happen; they add zero. zip(..., strict=True) pairs coordinates and raises ValueError if the lengths differ. max(..., key=lambda i: scores[i]) returns the index with the largest score. A dense matrix multiply performs the same $n \times d$ coordinate products, often much faster than a Python loop.

Even though only coordinate 0 of the question vector is nonzero, this dense loop still performs all 12 multiplications (4 chunks $\times$ 3 coordinates). A dense matrix operation doesn't automatically exploit arbitrary zero values. Sparse representations and specialized kernels can change the work, but they require a different execution path.

With 50,000 chunks and 768 coordinates, the same pass uses 38,400,000 multiplications before it chooses citations. The exact score pass is Θ(n * d) in time.

Storage has its own ledger. Those float32 vectors take 50,000 * 768 * 4 = 153,600,000 bytes, about 146.5 MiB: Θ(n * d) input memory, separate from temporary workspace. The build-it function later materializes its allowed rows, adding O(n) temporary references; a streaming exact scan can keep only the query vector and running winners in extra memory.

Return top-k without sorting every result

The assistant will cite only k = 2 chunks. A first version can score all candidates and fully sort every result. Before optimizing, notice the mismatch: the system needs two winners, but the sort orders every loser too. In product code, treat k <= 0 as an empty result so heap logic never reads from an empty list.

This is correct. It also orders two losers that will never be shown. Full sorting costs O(n log n) comparisons in the worst case and materializes all n scored results. A request that needs only a small k shouldn't pay to order every loser. What smaller structure can keep the weakest winner easy to replace?

In Data Structures for AI you already used a min-heap so the weakest retained result sits at the root. A bounded min-heap stores only the top-k values seen so far. A stronger candidate can replace the current weakest winner. Repairing the heap follows one path through a tree with k entries. That path has at most about log2(k) parent-child steps: three for k = 8, ten for k = 1,024.

Python's heapq helpers implement that structure in the lab. Tuples compare left to right, so (score, id) breaks a score tie using the ID: this code favors the lexicographically larger string. Change the key if your requirement differs; for integer IDs, (score, -doc_id) makes smaller IDs win ties.

For 2 <= k <= n, maintaining the heap costs O(n log k) after scoring, with O(k) selection memory. Sorting the final winners adds O(k log k), still within that bound. For k = 1, one running maximum gives O(n) selection; for k > n, there are only n items to retain.^{[5]Reference 5Introduction to Algorithms (3rd ed.)https://mitpress.mit.edu/9780262033848/introduction-to-algorithms/}

The next function accepts the same scored rows and keeps at most two winners. Compare its result with the full sort:

The heap improves selection, not exact vector scoring. You still paid O(n * d) to produce the four scores above, and you'd still pay 38.4 million coordinate products at n = 50,000, d = 768. A faster tail step can't fix an expensive upstream step.

Keep separate time and memory ledgers

The same algorithm can have one growth rate for work and another for storage. Before changing code, label each number as input memory, output memory, or temporary workspace:

A heap can cut selection memory from Θ(n) to Θ(k), but it can't shrink a vector matrix that the service keeps resident. Likewise, changing a nested loop to stream its counter can save memory while leaving its Θ(n²) time unchanged. Time and memory need separate fixes.

When an exact dot-product scan blows past your compute budget at scale, the algorithmic remedy isn't to sort faster: it's to change the candidate-search path by using an Approximate Nearest Neighbor (ANN) index such as Hierarchical Navigable Small World (HNSW).

Spot a quadratic trap

Linear work doubles when the input doubles. Quadratic work grows far faster because each item is compared with many other items. The quickest way to spot it: ask whether the inner loop's range grows with the outer loop.

A tempting cleanup step might compare every retrieved chunk against every other retrieved chunk to remove near-duplicates. For n chunks it evaluates n * (n - 1) / 2 pairs, which is Θ(n²). Before running the counter, estimate the jump from 16 to 1,000 chunks; the exact output makes the trap hard to ignore.

At n = 1,000, a request-time pairwise pass makes 499,500 comparisons. That's a concrete failure case: a harmless-looking loop can consume the latency budget before generation begins. The count came from the changing inner range, not from the fact that the code has two for statements.

Two nested loops aren't automatically quadratic. If the inner loop visits exactly 32 tokens for each chunk, the count is 32 * n, which is Θ(n) because 32 doesn't grow with n. Count the iterations, not the indentation.

If the requirement is only exact duplicate removal, a set changes the operation. It remembers texts already seen instead of comparing every pair. That is a change in the requirement, not a magic rewrite of semantic similarity.

This path performs exactly n membership checks. Under normal hash-table behavior, those checks take expected O(n) total time because each is expected O(1); collisions can make a single check O(n). Hashing still reads each text, so unbounded string length remains part of the real cost.

It doesn't solve semantic near-duplicate detection. Requirements determine the algorithm: don't silently replace "similar meaning" with "identical string" just to obtain a better complexity label.

The same square shape appears later in attention

The same counting habit will reappear in Transformer internals. For now, carry one preview: scaled dot-product attention computes a score from every query position against every key position. Vaswani et al. write that matrix as $QK^{\top}$, so a self-attention sequence with t token positions has t * t score cells.^{[6]Reference 6Attention Is All You Need.https://arxiv.org/abs/1706.03762}

You don't need the rest of the model yet. If the prompt doubles from 256 to 512 tokens, predict the cell count before looking at the program.

Doubling tokens quadruples cells. Treat this as a per-layer, per-head order of magnitude for a dense float32 score grid for one sequence during prefill, before multi-head batching, fused kernels, or KV-cache decode optimizations.

If each query and key has $d_k$ coordinates (the per-head width in later chapters), forming that grid also performs $\Theta(t^{2} d_{k})$ multiply-accumulate work per head, while the score grid itself takes $\Theta(t^{2})$ values. This doesn't model a complete language model or serving memory yet. It transfers the algorithm skill: when you see an all-pairs matrix, ask whether an input is being squared.

Prefill versus decode: two computational regimes

An autoregressive Transformer has two useful phases to count. The dense-attention costs below are per head, with head width $d_k$; they don't include projections, feed-forward layers, or the rest of the serving request:

1. Prefill (prompt evaluation): Process the prompt's $T$ positions together, possibly in chunks. Causal attention lets each position attend to itself and earlier positions. Dense attention work is $\Theta(T^2 d_k)$ per head, up to the causal constant factor. Matrix-matrix operations can reuse loaded values well, so large prefills often become compute-bound. Prefill contributes to time to first token (TTFT), which also includes queueing, preprocessing, and other work.
2. Decode (token generation): For one sequence, a step adds one query against a cached history of length $t$. Its attention work is $\Theta(t d_k)$ per head. Small decode batches often become bandwidth-bound while reading model weights and KV state. Larger batches, kernel choices, and hardware can change the bottleneck. Decode contributes to inter-token latency (ITL); batching and scheduling also affect the gaps observed by a client.

Neither phase has one universal FLOPs-per-byte value. A batched decode also combines work across sequences, so it needn't execute as one literal vector-matrix kernel. Fused attention such as FlashAttention can avoid storing the full score grid in device memory while retaining dense attention's quadratic arithmetic.^{[7]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

Attention compute scaling versus KV cache memory

The KV cache stores past key and value activations so decode doesn't recompute them at every step. One dense-attention decode step is linear in cached length, but the cache grows as generation continues. Across $G$ generated tokens after a prompt of length $T$, attending to the growing history totals $\Theta((GT + G^2)d_k)$ work per head. Caching saves recomputation; it doesn't make an arbitrarily long continuation linear in total output length.

For Llama 3 8B's 32 layers, 8 Grouped-Query Attention (GQA) key-value heads, and head dimension $d_k = 128$, assume a 16-bit KV cache (2 bytes per element).^{[8]Reference 8Llama model specificationshttps://github.com/meta-llama/llama-models/blob/main/models/sku_list.py} Each stored token then requires:

$$\text{KV memory per token} = 2 \times 32 \text{ layers} \times 8 \text{ KV heads} \times 128 \times 2 \text{ bytes} = 131{,}072 \text{ bytes} = 128 \text{ KiB}.$$

At a concurrent batch size of $B = 16$ streams, context length directly dictates your serving memory budget:

• At 1,024 tokens per stream: $16 \times 1{,}024 \times 128 \text{ KiB} = 2.0 \text{ GiB}$.
• At 4,096 tokens per stream: $16 \times 4{,}096 \times 128 \text{ KiB} = 8.0 \text{ GiB}$.
• At 8,192 tokens per stream: $16 \times 8{,}192 \times 128 \text{ KiB} = 16.0 \text{ GiB}$.

These counts include raw KV values only, with equal-length streams and no prefix sharing. Model weights, temporary workspace, cache metadata, and unused space in partially filled blocks consume additional memory. Cache quantization or sharing changes the count.

Measure scale and install a budget guard

Big-O tells you the shape, not whether today's workload meets a service objective. Count expected work to spot dangerous growth, then benchmark the real implementation on representative data and hardware. A guard can reject or reroute requests whose predicted work exceeds the supported budget.

The first guard uses a multiply count, not a made-up millisecond conversion. That keeps the example honest: hardware and kernels determine how many milliseconds those operations take.

The rejected case doesn't justify weakening authorization or returning arbitrary evidence. Exact ranking over that many authorized candidates no longer fits this budget. Possible next experiments are a stricter legitimate filter, an offline-created index, or an approximate nearest-neighbor index evaluated for recall.

A later option is HNSW, a hierarchical proximity-graph approach to approximate nearest-neighbor search. It trades exact scoring of every vector for graph-guided candidate exploration. The original paper reports logarithmic scaling under its construction assumptions; treat that as a reason to measure recall and latency together, not as a free O(log n) replacement for the exact baseline.^{[9]Reference 9Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphshttps://arxiv.org/abs/1603.09320} That belongs in Vector DB Internals, not in this foundation exercise.

Guarding tail latency in production

Counting algorithmic operations gives you the asymptotic shape, but enforcing wall-clock budgets in production requires tracking where runtime time disappears:

• Look beyond the median: The mean and p50 are different statistics, and neither describes the slowest requests. An invented service could have a 20ms median and an 800ms p99. Measure the percentiles and deadlines your users need, then include queue time as well as computation in the budget.
• Profile Python collection: CPython uses reference counting plus a cyclic garbage collector. Collection can contribute pauses, but their duration depends on the object graph and runtime. Measure collection activity before changing it. gc.disable() stops automatic cyclic collection; reference counting continues, and unreachable cycles can accumulate without a deliberate collection policy.^{[10]Reference 10gc: Garbage Collector interfacehttps://docs.python.org/3/library/gc.html}
• Bound batch formation time: A rule that waits for exactly eight arrivals can delay a low-traffic request indefinitely. Ray Serve's batch_wait_timeout_s limits how long it tries to fill a batch after the first request; Triton has its own max_queue_delay_microseconds setting. These settings don't cap total request latency when execution is already busy.^{[11]Reference 11Ray Serve: Dynamic Request Batchinghttps://docs.ray.io/en/latest/serve/advanced-guides/dyn-req-batch.html[12]Reference 12Triton Inference Server: Dynamic Batcherhttps://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/batcher.html}
• Measure preprocessing separately: Tokenization or cleaning can leave an accelerator waiting for input. Worker pools and overlapping preparation with execution can help, but they add scheduling, memory, and transfer costs. Compare the measured pipeline before and after rather than assuming more workers keep the GPU busy.

A separate architecture is retrieve, then refine: an inexpensive first-stage score proposes a shortlist, then a slower scorer or cross-encoder ranks only those survivors. A one-stage top-k heap isn't this pattern: it keeps the best rows under one score and returns those same scores. Refinement requires a second scoring function, and the shortlist must preserve authorization and enough recall for the final answer.

The fixture below makes that second stage visible using precomputed scores, not real model calls. Its rows have already passed the permission filter. The first stage reads four cheap scores; the second reads richer scores for only three shortlisted rows. A real reranker would compute those richer scores on demand. Here the counters track score reads so you can inspect which rows reach each stage.

This is two-stage retrieval: the cheap score chooses candidates, and the richer score changes their order. If the first stage drops a relevant row, the second stage can't recover it, so measure shortlist recall before celebrating lower refinement cost.

Build it: an authorized exact top-k retriever

Now assemble an exact baseline with a guard that runs before vector scoring. This permission fixture contains three allowed auth rows, one private auth row, and one billing row. Only the allowed rows contribute to the scoring budget:

• Authorization happens first.
• Every authorized vector is scored exactly.
• A bounded heap retains only k citations.
• A preflight count rejects work above the configured multiply budget.

The first call scores three allowed vectors and returns two citations. The second call has room for only eight multiplies, so it fails before performing any dot product. Authorization still scans the input rows, and the guard doesn't bound that work or elapsed time. It protects one named stage.

For this lab, vectors contain finite numbers and must have the query's dimension. zip(..., strict=True) detects a mismatched dimension instead of silently dropping coordinates. Compare any optimized retriever against this exact ranking on the same allowed rows, with the same tie rule.

Check the growth by hand

Work these exercises without code first. Write down the growing input, the operation being counted, and the expected result before checking the sketches.

1. A project has n = 2,500 authorized chunks, each with d = 384 coordinates. How many coordinate multiplications does an exact dot-product scan perform?
2. A full sort orders 10,000 scored results even though the assistant returns k = 5. Which part does a bounded heap reduce, and which part remains unchanged?
3. A near-duplicate cleanup compares each of 200 chunks with every later chunk. How many comparisons occur?
4. An attention preview changes prompt length from 256 tokens to 1,024 tokens. By what factor does the number of score cells grow?
5. Your guard allows one million coordinate multiplications. Does exact scoring fit for n = 4,000, d = 256?
6. Three cached prompts share the first 6 tokens with a new 8-token request, then diverge. About how many token comparisons does a naive per-key scan do if each key stops at the first mismatch?
7. A sorted catalog has 65,536 chunk IDs. What is the worst-case number of midpoint checks for binary search?

Solution sketches

1. 2,500 * 384 = 960,000 coordinate multiplications.
2. The heap reduces selecting the best five values from full-sort work to bounded top-k work. It doesn't remove the exact score pass that produced all 10,000 values.
3. 200 * 199 / 2 = 19,900 comparisons.
4. Token length grows by 1,024 / 256 = 4, so square score-cell count grows by 4 * 4 = 16.
5. No. 4,000 * 256 = 1,024,000, which exceeds the one-million budget.
6. Each of the 3 keys compares 6 matches plus 1 mismatch, so about 21 comparisons. A radix-tree walk would scale with the prefix length, not with n.
7. At most 17 midpoint checks. Although $2^{16} = 65{,}536$, this half-open binary search can spend a seventeenth comparison checking the final one-item interval. For example, searching for the first ID visits 17 midpoints.

Which input grows which cost

The recap table puts one final label on each operation: which input grows, what shape follows, and what an engineer should question first.

Three rules carry into every later AI system chapter:

1. Name the operation and its growing inputs before stating a complexity.
2. Keep a smallest-correct baseline so optimization can be checked for correctness.
3. Couple latency work with product requirements: authorization, citation quality, and recall don't disappear when a faster algorithm arrives.