Capstone — Build a Document Search Engine by Judging AI Code

Across this course you have learned data structures, complexity analysis, and how to work alongside AI. This chapter ties it all together into one repeatable workflow, exercised on a single real project that touches every data structure you’ve studied — search, hash tables, sorting, and graphs. This is the workflow you’ll use to attack new problems for the rest of your career, no matter how good AI gets.

The six-step workflow

Being a professional problem-solver is not about “writing code fast.” It is about thinking systematically before you write. This workflow has six steps you can follow in order on any problem.

1. Decompose — Break the big problem into small, understandable pieces. State clearly “what is the input, what is the output, and what are the constraints” before thinking about a solution.

Key question: What is this problem really asking, and which sub-problems must I solve?

2. Choose structure — The structure you pick determines the speed of every operation. Ask whether you search often, need to keep order, or access by key versus by index.

Key question: Which operation happens most often, and which structure makes it fastest?

3. Estimate complexity — Before writing code, estimate the Big-O of your intended approach. If it comes out O(n²) on a million items, find out now — not when production falls over.

Key question: If the input grows 10×, by how much does the running time grow?

4. Implement (AI allowed) — Now write the code, and use AI freely. Because you already have a plan, AI becomes a “fast pair of hands,” not “the brain that decides for you.”

Key question: Does this code match the complexity plan I intended?

5. Verify — Test against normal cases, edge cases (empty list, duplicates, a single value), and confirm the real complexity matches your estimate.

Key question: What kind of input would make this code break or slow down unexpectedly?

6. Improve — Only once it is correct, ask whether it can be clearer, faster, or more concise. Optimizing before it is correct is wasted effort.

Key question: Where can I make this clearer or faster without breaking correctness?

When working with AI, these six steps collapse into four in practice: Prompt (give the task after you already have a plan), Read (read the AI’s code carefully), Judge (rule on correctness and complexity), Improve (fix what it missed). We’ll run this cycle four times on one project below.

The capstone project: a search-and-recommendation engine

Suppose you must build a system for a digital document library. Users need to (1) type a query and see autocomplete suggestions, (2) search documents without duplicates polluting the results, (3) see results ranked by relevance, and (4) see documents “related” to the one they’re currently reading. This is one system requiring four data structures working together:

Stage	Job	Core structure	Search/ordering technique
1	Autocomplete	Sorted list	Binary search
2	Deduplicate + index	Hash table	Hash as key
3	Rank results	List of pairs (score)	Sorting
4	Suggest related docs	Graph (adjacency list)	BFS

We’ll walk all four stages with the same cycle every time: Prompt → Read → Judge → Improve. AI will give a “correct but unscaled” answer every time — your job is to catch it before it reaches production.

Stage 1 — Search: autocomplete

Prompt sent to AI: “Write a function search_prefix(words, prefix) that returns every word in a word list that starts with prefix. This system will be called thousands of times per second as the user types.”

AI’s answer (looks correct at a glance):

def search_prefix(words, prefix):
    """Return every word that starts with prefix."""
    return [w for w in words if w.startswith(prefix)]

Read: The code is 100% correct. It passes every small test case. It’s readable, with no logic bugs at all.

Judge: The problem isn’t correctness — it’s complexity under repeated calls. Each call costs O(n · L) (n = number of words, L = average word length). An autocomplete system calls this function on every keystroke. With q calls per second, that’s O(q · n · L) total — with 200,000 words and a fast typist, the lag becomes noticeable immediately.

Improve: Sort the word list once when the system is built, then use binary search to find the boundary of matching words, instead of scanning the whole list every single call.

(The widget above demonstrates the binary search principle on sorted numbers. The same principle applies to searching words in an alphabetically sorted list.)

import bisect

class PrefixSearcher:
    """Sort the word list once, then search prefixes repeatedly and fast."""

    def __init__(self, words):
        self.words = sorted(words)              # O(n log n), once at construction

    def search(self, prefix):
        lo = bisect.bisect_left(self.words, prefix)
        hi = bisect.bisect_left(self.words, prefix + "￿")  # upper bound: a char greater than any
        return self.words[lo:hi]                # O(log n + k), k = number of matches

# Verify
s = PrefixSearcher(["cat", "car", "dog", "care", "card"])
assert s.search("car") == ["car", "card", "care"]
assert s.search("z") == []                      # no match
assert s.search("") == sorted(["cat","car","dog","care","card"])  # empty prefix = everything

The preparation cost O(n log n) happens once. Every search after that costs only O(log n + k) — a dramatic speedup once there are many repeated queries. This is the classic pattern: pay an upfront cost once, in exchange for speed forever.

Stage 2 — Hash: deduplicate before indexing

Prompt sent to AI: “Write a function find_duplicate_docs(docs) that returns groups of document indices whose content matches exactly, character for character.”

AI’s answer (correct, but hiding a problem):

def find_duplicate_docs(docs):
    groups = []
    used = set()
    for i in range(len(docs)):              # O(n)
        if i in used:
            continue
        group = [i]
        for j in range(i + 1, len(docs)):   # O(n) nested inside!
            if docs[j].strip() == docs[i].strip():
                group.append(j)
                used.add(j)
        if len(group) > 1:
            groups.append(group)
    return groups

Read: The result is correct in every case. Small test examples pass. The structure is easy to follow.

Judge: There are two nested loops — comparing every pair of documents is O(n²), and each comparison also compares characters, O(L). Total: O(n² · L). A library with 100,000 documents would require roughly 10 billion comparisons — unusable in practice.

Improve: Use a hash table (dict) to group by content instead of comparing every pair — identical documents fall into the same key automatically.

from collections import defaultdict

def find_duplicate_docs(docs):
    """Return groups of document indices with identical content (duplicates only)."""
    groups = defaultdict(list)          # key(content) -> list of indices
    for i, text in enumerate(docs):     # O(n)
        key = text.strip()              # normalize before comparing
        groups[key].append(i)           # O(1) average
    return [idxs for idxs in groups.values() if len(idxs) > 1]

# Verify (with edge cases)
assert find_duplicate_docs([]) == []                        # empty list
assert find_duplicate_docs(["a"]) == []                     # single item, no dup
assert find_duplicate_docs(["a", "a", "a"]) == [[0, 1, 2]]   # three duplicates
assert find_duplicate_docs(["a ", "a"]) == [[0, 1]]          # trailing space -> strip makes them equal

We walk documents once, O(n) of them; each computes a key proportional to its length, O(L), giving O(n · L) total — far better than the all-pairs O(n² · L).

Extending it — a full-text inverted index: With duplicates gone, we use the same hashing principle to build an “inverted index” — a dict keyed by word, whose value is the set of document ids containing that word.

def build_inverted_index(docs):
    index = defaultdict(set)
    for doc_id, text in enumerate(docs):        # O(n)
        for word in set(text.lower().split()):  # O(L) per document
            index[word].add(doc_id)             # O(1) average
    return index

# Find documents containing "python" -> O(1) average, instead of scanning every doc every time

Building the index costs O(total words across all documents) once. After that, each word lookup is O(1) average — the same pattern as Stage 1: pay once upfront, get speed forever.

Stage 3 — Sorting: rank the search results

Prompt sent to AI: “Sort the search results (pairs of document id and relevance score) from highest score to lowest.”

AI’s answer (works correctly, but…):

def sort_by_score(results):  # results: [(doc_id, score), ...]
    n = len(results)
    for i in range(n):
        for j in range(n - 1):
            if results[j][1] < results[j + 1][1]:
                results[j], results[j + 1] = results[j + 1], results[j]
    return results

Read: This is bubble sort, correctly sorting descending as requested. It works correctly on every small test.

Judge: Bubble sort is O(n²), always, no matter how nearly sorted the data already is. A real search system might need to sort tens of thousands of results per query — the AI wrote its own sorting algorithm when it wasn’t necessary, when the language already ships a well-tuned one for free.

Improve: Use the language’s built-in sort (Timsort in Python), which guarantees O(n log n).

def sort_by_score(results):
    return sorted(results, key=lambda r: r[1], reverse=True)   # O(n log n)

# If you only need the top k (typical for a results page), heapq is even faster
import heapq
def top_k_by_score(results, k):
    return heapq.nlargest(k, results, key=lambda r: r[1])      # O(n log k)

Key lesson: Don’t let AI write a standard algorithm (sorting, searching) from scratch unless there’s a genuinely specific reason (e.g. stability under special conditions, or a learning exercise). Standard libraries have been tested and tuned for decades.

Stage 4 — Graph: suggest related documents

Prompt sent to AI: “Two documents are connected if they share a tag. Write a function related_docs(graph, start, max_hops=2) that finds every document reachable within 2 levels (hops).”

AI’s answer (a more dangerous bug than the ones so far):

def related_docs(graph, start, max_hops=2):
    result = []
    def dfs(node, hops):
        if hops > max_hops:
            return
        for neighbor in graph[node]:
            result.append(neighbor)
            dfs(neighbor, hops + 1)     # no visited set!
    dfs(start, 0)
    return result

Read: The code loops over neighbors and calls itself recursively (DFS). At a glance it seems to cover every document within 2 levels.

Judge: This isn’t just a complexity issue — it’s a bug that crashes the program. There is no visited set at all. If the graph has a cycle (document A connects to B, and B connects back to A — which always happens, since a shared tag is bidirectional), the function recurses forever until a RecursionError. Even without a cycle, the result still contains the same document multiple times, reached via different paths.

Improve: Use BFS with a visited set — it guarantees each document is visited exactly once, and gives more precise control over hop count than DFS.

from collections import deque

def related_docs(graph, start, max_hops=2):
    visited = {start}
    result = []
    queue = deque([(start, 0)])
    while queue:                            # O(V + E) total
        node, hops = queue.popleft()
        if hops == max_hops:
            continue
        for neighbor in graph[node]:
            if neighbor not in visited:
                visited.add(neighbor)       # prevents revisits / infinite loops
                result.append(neighbor)
                queue.append((neighbor, hops + 1))
    return result

# Verify, including an edge case with a cycle
graph = {"A": ["B", "C"], "B": ["A", "D"], "C": ["A"], "D": ["B"]}
assert related_docs(graph, "A", max_hops=1) == ["B", "C"]
assert set(related_docs(graph, "A", max_hops=2)) == {"B", "C", "D"}
assert related_docs(graph, "A", max_hops=0) == []

The complexity is O(V + E) over the documents and connections actually visited — the missing-visited-set bug is one of the most common patterns when AI writes graph traversal code. Watch for it especially closely.

Wiring it together into one system

Once all four stages have passed the Prompt → Read → Judge → Improve cycle, we assemble them into a single class — each method is a verified stage:

class DocumentSearchEngine:
    def __init__(self, docs):
        clean_ids = [g[0] for g in find_duplicate_docs(docs)] or list(range(len(docs)))
        self.docs = docs
        self.index = build_inverted_index(docs)          # Stage 2: O(total words)
        self.autocomplete = PrefixSearcher(self._all_words(docs))  # Stage 1: O(n log n)
        self.tag_graph = build_tag_graph(docs)            # Stage 4: O(n · avg tags)

    def suggest(self, prefix):
        return self.autocomplete.search(prefix)           # O(log n + k)

    def search(self, query_words):
        matches = set.intersection(*(self.index[w] for w in query_words))
        scored = [(doc_id, self._score(doc_id, query_words)) for doc_id in matches]
        return sort_by_score(scored)                      # O(n log n)

    def related(self, doc_id):
        return related_docs(self.tag_graph, doc_id, max_hops=2)  # O(V + E)

Notice that each stage has its own complexity, and none of them makes another slower. This is why you estimate Big-O separately, per part before wiring the system together — if you wire first and measure speed afterward, you won’t be able to tell which part is the bottleneck.

Choosing & justifying

What separates a beginner from a professional is being able to explain why you chose that structure. Here is the reasoning across all four stages of this project:

Stage	Structure chosen	Rejected alternative	Trade-off
1. Search	Sorted list + binary search	Scan the whole list every time (O(n) per query)	Must re-sort on insertion (O(n) per insert)
2. Hash	dict/defaultdict	All-pairs comparison, O(n²)	Extra memory for keys; O(1) is average-case only
3. Sorting	sorted() (Timsort)	Hand-written bubble/selection sort	None — the standard library wins on every axis in the general case
4. Graph	BFS + visited set	DFS without visited (breaks on cycles)	BFS needs a queue, but is safe and gives precise hop control

A good justification does not just say “what you chose” — it says “why it beats the alternatives, and what you gave up in return.”

Judge’s Checklist

This is the checklist you saw applied four times above — memorize it, and use it to judge any AI code for the rest of your career, no matter how smart AI gets.

• Correctness: Does it give the right result on normal input?
• Edge cases: Have you tested empty list, a single value, all-duplicates, negatives, a graph with a cycle, and very large input?
• Complexity: What is the real Big-O, and does it match what you intended? Any hidden nested loops or in on a list?
• Data structure: Is it the right one for the job? Is there a structure that would make it faster?
• Readability: Are variable names meaningful, the structure clear, and comments present where the logic is tricky?
• Robustness: Does it handle invalid input gracefully?
• Common tells: Is there an in on a list inside a loop? A nested loop hidden inside a method? Graph/tree traversal with no visited set? A hand-rolled standard algorithm (sort/search) when the library already has one?

Memorize four questions: Is it correct? Do the edges break it? How complex is it, really? Is it readable? These four questions catch almost every mistake AI makes.

Open challenges (your turn — extend each stage)

These challenges extend all four stages of the project. There is no ready-made answer here — think it through yourself first, and only open the hint if you’re truly stuck.

A. Typo-tolerant autocomplete — Stage 1 uses exact character-by-character matching, but users mistype often (“pyhton” should still find “python”). How would you design the data structure and algorithm to stay fast enough for every keystroke?

Hint

Consider the edit distance between the prefix and each candidate word. But computing edit distance against every word is O(n) per query — look for a structure that narrows the search space before computing edit distance at all, such as a trie that prunes quickly, or an n-gram index that filters out words with no chance of being close.

B. Detect similar documents, not just identical ones — Stage 2 only detects content that matches character for character. How would you extend it to catch documents that differ only in whitespace, letter case, or even minor word reordering, while staying close to O(n)?

Hint

The first step is to normalize content before hashing (lowercase it, collapse repeated whitespace, strip punctuation). That alone handles the whitespace/case cases with the same hashing approach. For the “similar but not exact” case (reordered words), look into how MinHash or shingling solves this while still keeping hashing as the core mechanism.

C. Rank results with several scores at once — Stage 3 sorts on a single score, but a real system must blend relevance with recency, and needs a stable tie-breaking rule when scores are equal. How would you design the sort key?

Hint

Python’s sorted() is already a stable sort (it preserves the original order of equal-scoring items) — take advantage of this by passing a key that’s a tuple of criteria in priority order, e.g. key=lambda r: (-r.relevance, -r.recency). Think through why the order of the tuple’s elements matters.

D. The “most related” path, not just the “nearest” — Stage 4 uses BFS, which counts every edge as one hop regardless of how strong the connection is. But what if some tag pairs are “more related” than others (different weights)? Does BFS still give the best answer?

Hint

BFS only guarantees the shortest path when every edge has equal weight. Once weights differ, you need an algorithm that picks the path with the “lowest total cost,” not the “fewest steps.” Look up Dijkstra’s algorithm, and see why it needs a priority queue instead of a plain queue.

Self-assessment rubric

Weight	Component	What it measures
40%	AI Code Critique	Your ability to read AI’s code, judge correctness and complexity, spot flaws, and improve it yourself.
30%	Project w/ justification	Quality of a genuinely working solution, clarity of your structure-choice reasoning, and correctness of the Big-O analysis.
20%	Concept quizzes	Understanding of fundamentals — data structures, complexity, and when to use which.
10%	Participation	Doing the exercises, asking questions, and joining discussion throughout the course.

Exercises

1. Two Sum — end-to-end. Walk the six steps briefly, then write a function that takes a list nums and a target, returning the indices of two numbers that sum to target.

Answer

Decompose: find a pair summing to target — Choose: dict remembering seen values — Estimate: O(n) — Implement — Verify edge cases — Improve: one pass is enough.

def two_sum(nums, target):
    seen = {}                       # value -> index
    for i, x in enumerate(nums):    # O(n)
        need = target - x
        if need in seen:            # O(1)
            return [seen[need], i]
        seen[x] = i
    return None
# two_sum([2, 7, 11, 15], 9) -> [0, 1]

Instead of comparing every pair O(n²), we remember past values in a dict and ask whether target - x was seen in O(1), for O(n) total.

2. Count frequencies and find the most common. Given a list of items, return the item that appears most often. Walk all six steps.

Answer

Choose: dict/Counter because we update counts repeatedly — Estimate: O(n)

from collections import Counter

def most_common_item(items):
    if not items:                   # edge case: empty list
        return None
    counts = Counter(items)         # O(n)
    return counts.most_common(1)[0][0]
# most_common_item(["a","b","a","c","a"]) -> "a"

Counting with a list and per-item search would be O(n²), but a dict makes each update O(1), for O(n) total.

3. Balanced parentheses. Given a string of brackets ()[]{}, check whether they are correctly matched.

Answer

Choose: a stack, because the bracket opened last must close first (LIFO) — Estimate: O(n)

def is_balanced(s):
    pairs = {")": "(", "]": "[", "}": "{"}
    stack = []
    for ch in s:                    # O(n)
        if ch in "([{":
            stack.append(ch)
        elif ch in pairs:
            if not stack or stack.pop() != pairs[ch]:
                return False
    return not stack                # no leftover open brackets
# is_balanced("({[]})") -> True ; is_balanced("(]") -> False

Edge cases: empty string -> True, unclosed open -> False, extra close -> False.

4. Spot the bug in this binary search. This code intends to be an iterative binary search, but it hides one bug. Find it, then explain what input would make it fail.

def binary_search(arr, target):
    lo, hi = 0, len(arr)
    while lo < hi:
        mid = (lo + hi) // 2
        if arr[mid] == target:
            return mid
        elif arr[mid] < target:
            lo = mid
        else:
            hi = mid
    return -1

Answer

The bug is lo = mid, which should be lo = mid + 1 when arr[mid] < target. Without the +1, if lo == mid exactly (e.g. lo=1, hi=2, mid=1), the loop keeps re-computing the same mid forever — an infinite loop, not just a wrong result.

def binary_search(arr, target):
    lo, hi = 0, len(arr)
    while lo < hi:
        mid = (lo + hi) // 2
        if arr[mid] == target:
            return mid
        elif arr[mid] < target:
            lo = mid + 1         # fixed
        else:
            hi = mid
    return -1

Verify it against a two-element list, e.g. binary_search([1, 2], 2) — that’s the edge case that catches this class of bug best.

5. Predict the output: BFS vs DFS. Given the graph {"A": ["B", "C"], "B": ["D"], "C": ["D"], "D": []}, predict the visitation order of BFS from “A” versus DFS from “A”, and explain why they differ.

Answer

BFS (uses a queue, level by level): A, B, C, D — visits B and C (level 1) before D (level 2).

DFS (uses a stack/recursion, goes deep first): A, B, D, C — descends to B, then deep to D, before backtracking to C.

Both visit every node in the same O(V + E); they only differ in order. BFS suits “find the shortest path”; DFS suits “explore as deep as possible first,” e.g. checking connectivity or finding cycles.

6. What’s the real Big-O? This code builds an inverted index. State the Big-O in terms of n (number of documents) and L (average words per document), and explain your reasoning.

def build_inverted_index(docs):
    index = {}
    for doc_id, text in enumerate(docs):
        for word in text.split():
            if word not in index:
                index[word] = []
            if doc_id not in index[word]:      # <- watch this line
                index[word].append(doc_id)
    return index

Answer

At a glance it looks like O(n · L), since it’s just a “documents × words per document” nested loop. But the line if doc_id not in index[word] searches a list, which is O(k) where k is the number of documents already recorded for that word. In the worst case (a word appearing in nearly every document), that becomes O(n) per check, making the whole thing O(n² · L) in the worst case.

The fix: change index[word] from a list to a set — if doc_id not in index[word] drops to O(1) average, restoring the whole function to the intended O(n · L). This is the exact same pattern found in Stage 2 of the project — “a hidden list turns a hoped-for O(1) into O(n).”

AI Code Critique — the final boss

The integrative one: you ask AI to write a function that counts “the number of duplicate pairs” of items in a list. AI returns code that is correct but hides a complexity flaw — the same pattern found in every stage of the project above. Run the full Read → Judge → Improve loop again yourself, with no hints given first.

def count_duplicate_pairs(items):
    """Count how many pairs (i, j) with items[i] == items[j] and i < j."""
    seen = []
    count = 0
    for x in items:
        if x in seen:               # looks O(1) at a glance, but...
            count += seen.count(x)  # how many of this value already seen
        seen.append(x)
    return count

Read: The code walks items once, so it looks O(n). But seen is a list, so x in seen is O(n) and seen.count(x) is O(n) too.

Judge: The real complexity is O(n²), not the O(n) the code’s shape suggests. With a million items it is unusably slow. This is a subtle complexity flaw — “correct result, but unexpectedly slow.” Notice this is exactly the same pattern as Stage 2 and Exercise 6: a hidden list turns something that looks like O(1) into O(n).

Improve: Switch seen to a dict (or Counter) so both the lookup and the count are O(1).

Answer

from collections import defaultdict

def count_duplicate_pairs(items):
    counts = defaultdict(int)
    pairs = 0
    for x in items:                 # O(n)
        pairs += counts[x]          # every earlier equal element pairs with this one -> O(1)
        counts[x] += 1              # O(1)
    return pairs
# count_duplicate_pairs([1, 2, 1, 1]) -> 3  (three pairs of the value 1)

A dict tracks how many of each value we have seen. Each time we meet a repeated x, it pairs exactly with all counts[x] earlier copies. Every operation is O(1), so the whole task is O(n) — fixing the O(n²) flaw the AI left behind. If you caught this pattern in three different places (Stage 2, Exercise 6, and here), the judge’s checklist has become instinct.

🎮 Game Dev: the data layer of a small roguelike

Everything in this course — hashing, sorting, trees, graphs — shows up together inside a single roguelike’s turn loop. Every system a roguelike needs (colliding entities, finding paths, picking loot, deciding who acts next) is really just “which data structure, which Big-O” wearing a fantasy skin. Below we sketch the whole data layer in Python, name the structure/algorithm behind each system, and run this chapter’s Prompt → Read → Judge → Improve lens on one naive subsystem.

System	Day	Structure / algorithm	Big-O
Collision detection	1–2	Spatial-hash grid (dict of cell -> entities)	O(1) average per query vs. naive O(n²)
Entity registry + inventory	2	dict keyed by entity id; inventory = list per entity	O(1) lookup/add
Enemy pathfinding	3	A* over the tile grid (priority queue + heuristic)	O(E log V)
Loot route planning	3	Grid DP (best score reachable per tile)	O(rows · cols)
Turn order	2	Priority queue (min-heap) keyed by “next action time”	O(log n) per turn

<rect x="150" y="20" width="130" height="60" rx="6" fill="none" stroke="#22c55e"/>
<text x="160" y="40" fill="currentColor">Entities+Inv</text>
<text x="160" y="58" fill="#22c55e">dict · O(1)</text>

<rect x="290" y="20" width="130" height="60" rx="6" fill="none" stroke="#38bdf8"/>
<text x="300" y="40" fill="currentColor">Enemy AI</text>
<text x="300" y="58" fill="#38bdf8">A* · O(E log V)</text>

<rect x="430" y="20" width="130" height="60" rx="6" fill="none" stroke="#f97316"/>
<text x="440" y="40" fill="currentColor">Loot Route</text>
<text x="440" y="58" fill="#f97316">Grid DP · O(r·c)</text>

<rect x="570" y="20" width="130" height="60" rx="6" fill="none" stroke="#ef4444"/>
<text x="580" y="40" fill="currentColor">Turn Order</text>
<text x="580" y="58" fill="#ef4444">Heap · O(log n)</text>

Game Loop one tick per turn

Figure: each roguelike system maps to one data structure and one Big-O bound; all five feed into a single per-turn game loop.

The naive subsystem, judged: collision detection

Prompt sent to AI: “Write a function nearby_entities(entities, x, y, radius) that returns every entity within radius tiles of (x, y). This runs once per frame for every monster on the level to check for collisions.”

AI’s answer:

def nearby_entities(entities, x, y, radius):
    """Scan every entity and keep the ones within radius."""
    return [e for e in entities
            if abs(e.x - x) <= radius and abs(e.y - y) <= radius]

Read: Correct. Passes every small test — 5 entities, 10 entities, all fine.

Judge: This is O(n) per call, and it’s called once per monster, per frame — O(n²) per frame across the whole level. A floor with 300 entities checking collisions 60 times a second is 5.4 million distance checks a second for this one query alone. Same shape of mistake as Stage 2 of the search engine above: correct code, unscaled complexity.

Improve: Bucket entities into a spatial hash — a dict keyed by (cell_x, cell_y) — so a query only scans the handful of entities in nearby cells instead of the whole level.

class SpatialHash:
    """Grid-bucketed entities: O(1) average insert, O(1) average query."""

    def __init__(self, cell_size: int = 4) -> None:
        self.cell_size = cell_size
        self.cells: dict[tuple[int, int], list] = {}

    def _key(self, x: float, y: float) -> tuple[int, int]:
        return (int(x // self.cell_size), int(y // self.cell_size))

    def insert(self, entity) -> None:
        self.cells.setdefault(self._key(entity.x, entity.y), []).append(entity)

    def nearby(self, x: float, y: float, radius: int = 1):
        cx, cy = self._key(x, y)
        found = []
        for dx in range(-radius, radius + 1):
            for dy in range(-radius, radius + 1):
                found.extend(self.cells.get((cx + dx, cy + dy), []))
        return found

# Verify
class Entity:
    def __init__(self, x, y):
        self.x, self.y = x, y

grid = SpatialHash(cell_size=4)
a, b, c = Entity(0, 0), Entity(1, 1), Entity(50, 50)
for e in (a, b, c):
    grid.insert(e)
assert set(grid.nearby(0, 0)) == {a, b}   # c is far away -> never scanned

Insert is O(1) average; a query only touches (2·radius+1)² cells, each holding a small, roughly-constant number of entities — O(1) average instead of O(n). Multiplied across every monster on the level, the whole frame drops from O(n²) to roughly O(n).

(Toggle between brute-force and grid mode above — watch the comparison counter collapse once entities are bucketed by cell.)

Entities & inventory: one dict, keyed by id

entities: dict[int, dict] = {}          # entity_id -> {"hp": ..., "pos": ...}
inventories: dict[int, list[str]] = {}  # entity_id -> item names

def pick_up(entity_id: int, item: str) -> None:
    inventories.setdefault(entity_id, []).append(item)   # O(1) amortized

# Verify
entities[1] = {"hp": 10, "pos": (0, 0)}
pick_up(1, "torch")
pick_up(1, "dagger")
assert inventories[1] == ["torch", "dagger"]
assert entities[1]["hp"] == 10          # O(1) lookup by id

Every lookup, damage tick, or inventory check is O(1) because entities live in a dict keyed by id — the same trick as the document-search engine’s inverted index, just with entities instead of words.

Enemy pathfinding: A*

import heapq

def astar(grid, start, goal):
    """grid[y][x] is True for a wall. Returns the shortest walkable path, or None."""
    def h(p):  # Manhattan distance heuristic
        return abs(p[0] - goal[0]) + abs(p[1] - goal[1])

    frontier = [(h(start), 0, start, [start])]
    seen = {start}
    while frontier:                                    # O(E log V) total
        _, cost, pos, path = heapq.heappop(frontier)
        if pos == goal:
            return path
        x, y = pos
        for nx, ny in ((x+1, y), (x-1, y), (x, y+1), (x, y-1)):
            if 0 <= ny < len(grid) and 0 <= nx < len(grid[0]) and not grid[ny][nx] and (nx, ny) not in seen:
                seen.add((nx, ny))
                heapq.heappush(frontier, (cost + 1 + h((nx, ny)), cost + 1, (nx, ny), path + [(nx, ny)]))
    return None

# Verify
grid = [
    [False, False, False],
    [False, True,  False],
    [False, False, False],
]
path = astar(grid, (0, 0), (2, 2))
assert path[0] == (0, 0) and path[-1] == (2, 2)
assert (1, 1) not in path               # wall tile is never entered

A* is BFS from Stage 4 of the search engine, upgraded with a priority queue and a heuristic so it explores promising tiles first — O(E log V) on the tile graph instead of exploring every tile equally.

(Click tiles to add walls, then run the search — the priority queue visibly prefers tiles closer to the goal instead of a plain BFS’s uniform spread.)

Loot route: grid DP

def best_loot_path(loot_grid: list[list[int]]) -> int:
    """Max loot collectible walking only right/down from top-left to bottom-right."""
    rows, cols = len(loot_grid), len(loot_grid[0])
    dp = [[0] * cols for _ in range(rows)]
    dp[0][0] = loot_grid[0][0]
    for r in range(rows):
        for c in range(cols):
            if r == 0 and c == 0:
                continue
            best = max(dp[r-1][c] if r > 0 else -1, dp[r][c-1] if c > 0 else -1)
            dp[r][c] = best + loot_grid[r][c]
    return dp[rows-1][cols-1]                          # O(rows · cols)

# Verify
loot = [
    [1, 3, 1],
    [1, 5, 1],
    [4, 2, 1],
]
assert best_loot_path(loot) == 12   # one optimal path: 1 -> 3 -> 5 -> 2 -> 1

Trying every right/down path by brute force is exponential (2^(rows+cols)); the DP table reuses each cell’s best-so-far exactly once, so the whole grid fills in O(rows · cols) — the same “cache the subproblem” idea as any DP table.

Turn order: a priority queue by speed

import heapq

class TurnQueue:
    """Min-heap ordered by next action time; faster entities act sooner."""

    def __init__(self) -> None:
        self.heap: list[tuple[float, int, int]] = []   # (next_time, entity_id, speed)

    def schedule(self, entity_id: int, speed: int, now: float = 0.0) -> None:
        heapq.heappush(self.heap, (now + 100 / speed, entity_id, speed))  # O(log n)

    def next_turn(self) -> int:
        next_time, entity_id, speed = heapq.heappop(self.heap)           # O(log n)
        self.schedule(entity_id, speed, now=next_time)                   # reschedule
        return entity_id

# Verify
q = TurnQueue()
q.schedule(1, speed=10)   # acts roughly every 10 ticks
q.schedule(2, speed=25)   # much faster -> acts almost every 4 ticks
assert q.next_turn() == 2   # faster entity goes first

A naive “sort every entity by speed every turn” costs O(n log n) per turn; a min-heap keeps only the next actor at the top and reschedules in O(log n) — the same list-vs-heap trade-off as Stage 3 of the search engine, just running once per turn instead of once per search.

(The widget re-sorts live as you change speeds — notice only the top of the heap needs to be right, not the whole ordering.)

Exercises

1. Collision, matched. A monster needs to check which other monsters are within melee range every frame. Which structure keeps that check fast as the level fills up, and what’s its average-case Big-O?

Answer

A spatial hash — a dict keyed by grid cell, mapping to the entities inside it. A query only scans the entity’s own cell plus its neighbors, so it costs O(1) on average instead of the naive O(n) scan over every entity.

2. Inventory, matched. Your inventory system stores entities in a plain list and looks up one by scanning for e in entities: if e.id == wanted: .... What’s the Big-O of one lookup, and how do you fix it without changing what data is stored?

Answer

O(n) per lookup — a full scan in the worst case. Fix it by keying a dict on entity id (entities[id] -> entity) instead of a list; the same data, but lookups drop to O(1) average.

3. Turn order, matched — harder. A prototype turn system runs entities.sort(key=lambda e: e.next_time) at the start of every turn and takes entities[0]. What’s the Big-O per turn versus a min-heap, and why does the gap matter once there are hundreds of entities acting thousands of times per battle?

Answer

Sorting the whole list every turn costs O(n log n) per turn — over t turns that’s O(t · n log n). A min-heap only needs O(log n) to pop the next actor and O(log n) to reschedule it, so t turns cost O(t log n) total. With 300 monsters and 5,000 turns per battle, that’s the difference between roughly 12 million and 2.7 million operations — and the heap doesn’t even need to touch the other 299 entities on each turn.

4. Trace the loot DP — harder still. Given the 2×2 loot grid [[2, 3], [1, 4]], trace best_loot_path by hand: fill in the DP table cell by cell, then state the final answer and the path that achieves it.

Answer

dp[0][0] = 2, dp[0][1] = 2 + 3 = 5, dp[1][0] = 2 + 1 = 3, dp[1][1] = max(dp[0][1], dp[1][0]) + 4 = max(5, 3) + 4 = 9. Final answer: 9, achieved by the path 2 -> 3 -> 4 (right, then down).

Challenge problems

A. Extend the roguelike — fleeing monsters. A* always paths toward a goal tile. A frightened monster should path away from the player instead — maximizing distance, not minimizing it. How would you adapt the systems above to give a monster a “flee” route, and what does it cost?

Approach

A* itself only knows how to search toward one fixed goal — it can’t directly search for “farthest tile.” Instead, run a bounded BFS/Dijkstra out from the player’s position up to some radius R (O(V + E) over the reachable tiles), pick the tile with the largest distance found within reach of the monster, then hand that tile to astar() as the goal exactly as before. Total cost stays a small constant multiple of one ordinary A* search — you’re just choosing a smarter goal tile first.

B. Profile & fix a bottleneck. Every monster on the level recomputes astar() toward the player’s current tile every single frame, even when neither the monster nor the player has moved. With 50 monsters at 60 fps, that’s 50 full O(E log V) searches per frame, dominating the frame budget. How would you confirm this is the bottleneck, and how would you fix it?

Approach

Confirm it first — don’t guess: wrap the per-monster update loop with cProfile or simple time.perf_counter() timestamps and check whether astar() calls account for most of the frame time. Once confirmed, cache each monster’s last path keyed by the player’s tile at the time it was computed, and only recompute when that tile actually changes:

path_cache: dict[int, tuple[tuple[int, int], list]] = {}

def get_path(monster_id: int, monster_pos, player_pos, grid) -> list:
    cached = path_cache.get(monster_id)
    if cached and cached[0] == player_pos:
        return cached[1]                      # player hasn't moved -> reuse
    path = astar(grid, monster_pos, player_pos)
    path_cache[monster_id] = (player_pos, path)
    return path

This turns “50 searches every frame” into “0 searches on most frames, 1 search per monster only when the player actually moves into a new tile” — the same instinct as caching in Stage 1 of the search engine: pay the expensive cost only when the underlying input has actually changed.

Going deeper & habits

Real course references:

• MIT 6.006 — Introduction to Algorithms (MIT OpenCourseWare) — full video lectures and problem sets, free.
• Harvard CS50 — Introduction to Computer Science (edX) — builds the foundation of computational thinking and problem-solving.
• Stanford CS161 — Design and Analysis of Algorithms — emphasizes complexity analysis and correctness proofs.

Books to read next — the shelf worth owning:

• CLRS (Cormen, Leiserson, Rivest, Stein), Introduction to Algorithms (4th ed.) — the field’s standard reference, covering every structure and technique used in this project with the most rigor and completeness.
• Skiena, The Algorithm Design Manual (3rd ed.) — best for “which technique when,” full of war stories from real systems, matching the systematic decision-making this project trains.
• Sedgewick & Wayne, Algorithms (4th ed.) — dives concretely into implementing data structures and sorting algorithms, extending Stage 2 and Stage 3 most directly.
• Kleinberg & Tardos, Algorithm Design — systematic design thinking for graph algorithms and dynamic programming, the best fit for extending Stage 4 and Challenge D.
• Bhargava, Grokking Algorithms (2nd ed.) — easy-to-follow illustrations, great for a fast refresher on everything above before tackling the more rigorous texts.

Habits for using AI without losing skill:

• Always plan first, then ask AI. Decompose, choose the structure, and estimate complexity yourself before asking. Use AI in the implement step, not the thinking step.
• Check the Big-O of every answer AI gives. Never assume “works” means “fast enough.” Hunt for hidden nested loops and in on lists every single time — all four stages of this project proved that this pattern recurs, whether in search, hashing, sorting, or graphs.
• Write it yourself; turn AI off periodically. Solve a problem by hand once a week to keep your problem-solving muscle alive. AI is an amplifier of skill, not a replacement for it.