🎓 A1Learner
🔍
0 of 112 lessons completed

DSA Intermediate Capstone — a Mini Search Engine

A search engine is a program that indexes documents by word and answers queries from that index. This capstone composes a hash map index, a trie and a heap in one program.

Build a tiny inverted index

Two short documents are split into words and counted. The index maps each chosen word to its per-document counts.

docs = [
    ("Garden Diary", "tomatoes basil tomatoes sun"),
    ("River Walks", "river stone river sun"),
]
index = {}
for title, text in docs:
    for word in text.split():
        per_doc = index.setdefault(word, {})
        per_doc[title] = per_doc.get(title, 0) + 1
for word in sorted(index):
    print(word, sorted(index[word].items()))

Trie suggestions for one prefix

A trie node holds a dict of next letters and an end flag. Words sharing a prefix share the start of their paths.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.end = False

class Trie:
    def __init__(self):
        self.root = TrieNode()
    def insert(self, word):
        node = self.root
        for ch in word:
            node = node.children.setdefault(ch, TrieNode())
        node.end = True
    def suggest(self, prefix):
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return []
            node = node.children[ch]
        found = []
        def walk(n, path):
            if n.end:
                found.append(path)
            for ch in sorted(n.children):
                walk(n.children[ch], path + ch)
        walk(node, prefix)
        return found

trie = Trie()
for word in ["garden", "garage", "gate", "river"]:
    trie.insert(word)
print("prefix ga:", trie.suggest("ga"))
print("prefix r:", trie.suggest("r"))

A heap ranks with a tie rule

Scores enter a heap as negative numbers beside their titles. Equal scores leave in alphabetical title order.

import heapq

scores = {"River Walks": 3, "Garden Diary": 3, "Night Sky Guide": 2}
heap = []
for title in sorted(scores):
    heapq.heappush(heap, (-scores[title], title))
while heap:
    neg, title = heapq.heappop(heap)
    print(title, "score", -neg)

Example

import heapq

class TrieNode:
    def __init__(self):
        self.children = {}
        self.end = False

class Trie:
    def __init__(self):
        self.root = TrieNode()
    def insert(self, word):
        node = self.root
        for ch in word:
            node = node.children.setdefault(ch, TrieNode())
        node.end = True
    def suggest(self, prefix):
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return []
            node = node.children[ch]
        found = []
        def walk(n, path):
            if n.end:
                found.append(path)
            for ch in sorted(n.children):
                walk(n.children[ch], path + ch)
        walk(node, prefix)
        return found

docs = [
    ("Garden Diary", "tomatoes and basil grow in the sunny garden the tomatoes ripen in the warm sun"),
    ("River Walks", "the river path winds past stone bridges the river runs cold in spring"),
    ("Kitchen Notes", "tomato soup needs basil garlic and slow heat a warm kitchen smells sweet"),
    ("Mountain Trails", "stone paths climb past cold springs the high trail crosses a windy ridge"),
    ("City Markets", "market stalls sell basil tomatoes and fresh bread every saturday morning"),
    ("Night Sky Guide", "bright stars wheel over the quiet ridge the moon climbs past midnight"),
]
index = {}
for title, text in docs:
    for word in text.split():
        per_doc = index.setdefault(word, {})
        per_doc[title] = per_doc.get(title, 0) + 1
trie = Trie()
for word in index:
    trie.insert(word)
print("documents indexed:", len(docs))
print("distinct words:", len(index))
for word in ["basil", "river", "stone", "tomatoes"]:
    print("index", word, sorted(index[word].items()))

def search(query):
    scores = {}
    for word in query.split():
        for title, count in index.get(word, {}).items():
            scores[title] = scores.get(title, 0) + count
    heap = []
    for title in sorted(scores):
        heapq.heappush(heap, (-scores[title], title))
    ranked = []
    while heap:
        neg, title = heapq.heappop(heap)
        ranked.append((title, -neg))
    return ranked

for query in ["tomatoes basil", "river stone", "the"]:
    print("search:", query)
    for place, (title, score) in enumerate(search(query), start=1):
        print(place, title, "score", score)
for prefix in ["ba", "st", "to"]:
    print("suggest", prefix, ":", trie.suggest(prefix))

Key points

  • A search engine indexes documents by word and answers queries from the index.
  • An inverted index is a map from each word to the documents holding it.
  • A document score is the total count of the query words inside it.
  • Ranked results print highest score first.
  • Ties in score break by document title in alphabetical order.
  • Prefix suggestions come from walking one path in a trie.
  • Suggestions print in sorted order, so every run agrees.

In this example, six short documents are indexed word by word. The index lines print each chosen word with sorted document counts. Each search adds the counts of its query words per document. A heap then ranks the scores, breaking ties by title. The trie answers three prefix queries in sorted order at the end.

Note

The corpus is six short original documents, written for this lesson.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Two documents tie on score for a query. Which one prints first?