Data Structures for AI

Data structures decide what an AI system can find, update, rank, and reuse quickly. If the structure matches the operation, a support bot can jump to the right help article, a background worker can process jobs in order, and a retrieval cache can reuse ranked document IDs without serving stale results.

Imagine a customer opens a support chat and types:

refund email

Behind the scenes, the company has thousands of help articles. The system can't read each article from start to finish under a tight latency budget. It needs a way to jump straight to the handful of articles that mention those words.

That jump is the job of a data structure. The NumPy lesson taught you that arrays and tensors are good when position is the whole story: if you know row 2, column 1, you can land on the value instantly. The CUDA and MPS chapters then showed where those tensors execute. This chapter answers the next systems question: what if your lookup doesn't start with a position? What if it starts with a word like refund, a repeated query like refund email, or a stream of jobs waiting to be processed?

Production AI products don't use one data structure for every job. A lexical retriever can index terms, an ingestion worker can queue jobs, and a retrieval cache can reuse ranked IDs only while their source documents remain current. Each structure exists because a particular operation repeats.

Picking a structure by the operation it must serve

Keep one question in mind: what operation must this structure serve?

• Do I need ordered items?
• Do I need exact lookup by key?
• Do I need uniqueness?
• Do I need the best few results?
• Do I need first-in, first-out (FIFO) work?
• Do I need to reuse previous results safely?

That is the real skill behind data structures. You are not memorizing names. You are matching an operation to a shape.

The previous NumPy lesson focused on position-based access. Here the query starts with a term, a top-k request, or a work queue, so the shape has to change.

One running example

We will use one tiny support-search story from start to finish. The documents are small enough to trace by hand, but the ideas scale to the retrieval layer inside a retrieval-augmented generation (RAG) system, where a large language model (LLM) answers with retrieved evidence.

User query:

refund email

We will use this story five times:

1. A list scan to show the smallest correct solution.
2. A dictionary of sets to make term lookup fast.
3. A heap to keep the top results.
4. A queue to show first-in, first-out work.
5. A cache to show safe reuse.

The list scan: the smallest working version

A list is the right starting point because the data is tiny and the order of documents is stable.

If you scan each document by hand, the match counts look like this:

Document 2 should rank above document 0.

This matters for beginners because a list scan is not "wrong." It is the smallest working version. The mistake is staying with a full scan after the question has changed.

When the corpus grows from 3 documents to 3 million, a query like refund email shouldn't force the system to reread each document. A full scan touches every document, so its cost grows in proportion to the corpus size. A dictionary lookup, by contrast, jumps to the value in roughly constant time on average regardless of how big the dictionary gets, because it hashes the key directly to its slot.^[1] That doesn't make the whole search constant-time: the retriever still reads the posting lists for the requested terms and ranks the resulting candidates. The next chapter formalizes this gap with Big-O notation; for now the intuition is enough: the data structure must let the query enter the system in its natural form, by term, not by document position.

Building an inverted index

The query starts with a term such as refund. That means the system should jump from a term to the matching documents.

This is what an inverted index does.

An inverted index maps each term to its posting list. A posting records a matching document ID; larger indexes can store additional information such as term positions or term frequency. In this first version, a Python set[int] is enough because our score asks only whether each query term appears in each document.^[2]

Build the index by hand

Split each document into terms and store the matching document IDs:

Now the query can enter through the index:

Turn those posting lists into match counts:

Document 2 wins because it matches both terms. Document 0 still matters because it matches one term.

Why a set lives inside the dictionary

The dictionary solves exact term lookup. The set solves uniqueness.

Suppose document 2 were refund refund email. The posting list for refund should still contain document 2 once, not twice. That is why the value is a set of document IDs instead of a list of document IDs.

This is a good beginner habit:

• Dictionary asks: "given this key, where is the value?"
• Set asks: "have I already seen this item?"

Production lexical indexes commonly keep ordered, compressed posting lists and richer statistics. The Python set is a teaching choice: it makes membership and duplicate prevention visible before you optimize storage or intersections.^[2]

Here is the smallest proof. The third support article says refund twice, but membership in the refund posting set still contributes one document ID.

Ranking with a heap

Once documents have scores, the system needs the best few results. It doesn't need a full sorted report about the corpus.

That is the moment for a heap.

For our tiny example, scores are:

A heap keeps the top k results easy to retrieve. In retrieval systems, that matters because you often want the best 5, 10, or 20 candidates, rather than a fully sorted candidate list.

For a stream of scored candidates, keep a min-heap capped at k items. The weakest retained score stays at the root. When a stronger candidate arrives, replace that weakest item instead of sorting every candidate seen so far.

Real rankers also need an explicit tie-break rule. The retriever lab below prefers the lower document ID when scores tie so its output stays deterministic.

The retriever pipeline

Here is how the pieces fit together in a single flow:

The query enters as plain text. Terms split apart. The index turns each term into a posting list. Scores accumulate per document. The heap returns the best k document IDs. This same pattern sits inside RAG systems, where a language model reads only the best supporting chunks instead of the whole corpus.

Code lab: a tiny retriever

Start with the slow version so you can see what the index is replacing. Then add the index and heap.

Put this in data_structures_demo.py. This first tokenizer lowercases and splits on whitespace. Production tokenizers also handle punctuation and language-specific rules, but keeping normalization tiny makes the structure visible. Returning a set means this example scores binary term membership: each distinct query term contributes at most once.

Read the data movement

The heap key uses score first and negative document ID second. Higher scores win; lower IDs win ties. That small rule makes repeated runs predictable.

That table is more important than the syntax.

What happens beyond retrieval: queues and caches

Retrieval is one place where data structures show up in AI systems.

Queues for background work

Imagine new documents arrive and need to be split into chunks, then encoded as embedding vectors before they can enter a semantic index. For one class of equal-priority jobs, arrival order is a clean baseline. Production queues may add priorities, retries, and dead-letter handling, but they still need an explicit ordering contract.

That is a queue.

The rule is simple:

• append new work on one side
• pop the oldest work from the other side

That is first-in, first-out behavior. If you accidentally process newest work first, users may see random delays or starvation in old jobs.

This failure case makes the ordering bug visible. New jobs belong at the tail; putting them at the front lets recent uploads jump ahead of old work.

Caches for repeated queries

If users ask the same question repeatedly, recomputing retrieval results per request wastes work.

A retrieval cache stores prior ranked document IDs so the system can reuse them.

Why the key includes version

If document 2 stops mentioning refund, the cached ranking [2, 0] becomes fast and wrong. A version number, timestamp, or explicit invalidation rule tells the system when old results must be dropped.

Run the failure before relying on the fix. A text-only key finds yesterday's ranking after an update; a versioned key misses and forces recomputation.

Common mistakes

Beginners often don't choose the wrong structure because they lack vocabulary. They choose it because the first working code feels good enough.

The fastest way to debug is still the same: trace one tiny example by hand.

Protect the promises with tests

These tests protect the structures, rather than the final answer alone:

Each test checks one promise:

• the index maps terms to the right documents
• empty input is predictable
• ranking follows visible score logic
• query normalization is consistent
• queue order stays first-in, first-out

Systems drill: simulate a stateful fulfillment lane

Production AI infrastructure is full of small systems with state, events, and invariants. Request queues, cache entries, retry state, tool side effects, and rollout gates all need explicit state transitions. If you can build one of these cleanly at small scale, you can transfer the pattern to larger LLM systems.

Here is a tiny fulfillment simulator. It processes events in order and refuses transitions that would break the system's promises.

Invariant

If you can name the invariant, write the state transition, and produce a failing case that proves the guard works, you can handle many "build a moderately complex system quickly" prompts. The same pattern applies to token schedulers, retry queues, tool-call ledgers, and evaluation runners.

Retries need a set too

Queue workers retry failed events. If a successful shipment event is delivered twice, applying it twice silently corrupts state. A completed-event ID set turns an accidental replay into a visible skip. Record the ID only after the side effect succeeds; a failed attempt must remain retryable.

That crash window is one reason the next SQL chapter matters. Durable rows, uniqueness constraints, and transactions let a system protect state across process failure.

Practice

Try these before you look at the solution sketches.

1. Add a fourth document: email support refund.
2. Predict sorted(index["refund"]).
3. Predict search("refund email").
4. Run search("invoice") and explain why it shouldn't crash.
5. Change document 2 so it no longer contains refund, then explain what must happen to the cache.

Solution sketches

Use these to check your reasoning, not to skip the exercise.

If your answer feels fuzzy, go back to the score table. Good data-structure reasoning should look inspectable on paper before it looks clever in code.

What to remember

Most beginner confusion disappears once you ask one concrete question:

What operation must stay fast and predictable?

Use that question to pick the structure:

This chapter intentionally stayed small. It doesn't try to cover all classic structures in one pass. That is a feature, not a missing piece. A beginner needs a handful of structures they can trace end to end before moving to bigger systems.

Later retrieval systems may combine exact term postings with approximate nearest-neighbor (ANN) indexes over embedding vectors. Hierarchical navigable small world (HNSW) is a layered graph index for ANN search; the Faiss work describes optimized GPU similarity search for brute-force, approximate, and compressed-domain setups. The core lesson doesn't change: store data in a structure that matches the lookup you need.^[3][4]

Three later patterns reuse the same ideas at larger scale:

• Beam search: a decoder needs the best few partial generations from a huge candidate set. A heap or top-k selection routine gives you candidates without sorting the full vocabulary at each step.
• Vector indexes: HNSW stores embeddings in a multi-layer graph so search can move from long-range links toward local neighbors. Its authors report logarithmic complexity scaling for their design, but ANN search still trades recall, latency, build time, and memory against one another.^[3]
• The key-value (KV) cache: during text generation, an LLM stores key and value vectors already computed for earlier tokens, then appends new entries as it decodes. It avoids recalculating those earlier key/value vectors, but the new query still attends over stored prefix state. The cache grows with sequence length and concurrent requests, which can limit how many requests a serving system can batch efficiently.^[5][6]

Handoff to the next chapter

Data structures decide which candidates enter the model's field of view.

That is why this chapter sits before databases. First you need to know which structure makes a lookup fast. Then you are ready to put those structures behind durable tables, joins, transactions, and vector search.

Mastery check

Key concepts

• lists
• dictionaries
• queues
• heaps
• indexes
• sets
• caches

Evaluation rubric

• Foundational: Explains what kind of operation each structure serves: ordered scan, exact lookup, uniqueness, top-k selection, FIFO work, or safe reuse
• Developing: Traces the support-search example by hand and predicts document scores before running code
• Intermediate: Builds the inverted index and shows why dictionary plus set is better than rescanning the full document list
• Applied: Uses a heap for top-k retrieval, a queue for background work, and a versioned cache for repeated queries
• Advanced: Writes tests for empty input, duplicate prevention, FIFO behavior, ranking logic, and cache invalidation after corpus changes, then explains why a durable write still needs idempotency

Follow-up questions

Common pitfalls

• A query slows down as the corpus grows. Cause: every request still scans the full list. Fix: build an index keyed by the field the query uses.
• Ranking counts feel too high or inconsistent. Cause: duplicates leaked into posting lists. Fix: store matching document IDs in sets when membership, not repetition, is the real state.
• Old embedding jobs never seem to finish. Cause: queue order was reversed by pushing or popping from the wrong side. Fix: test FIFO explicitly with a tiny trace.
• A cached ranking looks correct until the documents change. Cause: the cache key doesn't include version or invalidation state. Fix: tie cache entries to corpus updates.
• Heap output feels magical during debugging. Cause: scores stayed hidden inside loops. Fix: print one score table and compare heap output to the visible counts.
• A replay guard works in one process but duplicate writes still appear after crashes. Cause: completion marker and side effect are not one durable operation. Fix: enforce idempotency at the write boundary.