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 retriever can jump to the right model note, a background worker can process embedding jobs in order, and a retrieval cache can reuse ranked document IDs without serving stale results.

An engineer searches an internal model-serving notebook and types:

token hit

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

That jump is the job of a data structure. 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. CUDA and MPS explain where those tensors execute. Many systems start somewhere else: a word like token, a repeated query like token hit, 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

Ask one question before you pick a structure: what operation must it 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's the real skill behind data structures. You aren't memorizing names. You're matching an operation to a shape.

Position-based access works when coordinates are known. Here the query starts with a term, a top-k request, or a work queue, so the shape has to change.

A tiny retrieval story

One tiny technical-note search story carries the rest of the lesson. 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:

token hit

This story appears 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.

For beginners, a list scan isn't "wrong." It's 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 token hit 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, is expected amortized O(1) for the term key: it hashes the key directly to its slot.^{[1]Reference 1Designing Data-Intensive Applications.https://dataintensive.net/} That doesn't make the whole search constant-time. After the O(1) term lookup, scanning the posting list is O(|postings[term]|), so common terms whose postings are Θ(corpus) still dominate end-to-end lexical retrieval. The retriever then ranks the resulting candidates. The next chapter formalizes this gap with Big-O notation; for now, focus on the entry point: 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 token. 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]Reference 2Introduction to Information Retrieval.https://nlp.stanford.edu/IR-book/}

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 token token hit. The posting list for token should still contain document 2 once, not twice. That's 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]Reference 2Introduction to Information Retrieval.https://nlp.stanford.edu/IR-book/}

A repeated term is the smallest proof. In the demo below, document 2 says token twice, but membership in the token 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's 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 heap demo and the retriever lab below both prefer the lower document ID when scores tie so their output stays deterministic. The min-heap key stores score first and negative ID second, so higher scores win and lower IDs win ties when Python compares tuples.

The retriever pipeline

The pieces fit together in a single flow:

A query enters as plain text, splits into terms, and uses the index to turn each term into a posting list. Scores accumulate per document, and 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.

Building 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

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's a queue.

The rule is simple:

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

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

For a durable worker, FIFO order is only one promise. A stateful request lane should record each accepted event in a durable transaction, then preserve this invariant: every accepted job is pending, running, completed, or explicitly failed, never silently lost between memory and storage. The in-memory deque demonstrates ordering but provides no durability by itself.

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 token, 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

Checking your reasoning

Try these before you look at the solution sketches.

1. Add a fourth document: hit latency token.
2. Predict sorted(index["token"]).
3. Predict search("token hit").
4. Run search("attention") and explain why it shouldn't crash.
5. Change document 2 so it no longer contains token, 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.

Matching structures to operations

Most beginner confusion disappears once you ask one concrete question:

What operation must stay fast and predictable?

Use that question to pick the structure:

Later retrieval systems combine exact term postings with approximate nearest-neighbor indexes over embedding vectors. The selection rule stays the same: choose a structure that matches the lookup and update operations you need.

Moving to durable databases

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

That's why data structures come before databases. First you need to know which structure makes a lookup fast. Then you're ready to put those structures behind durable tables, joins, transactions, and vector search.