"""Reference solution — Inverted Index with Incremental Updates. DO NOT READ BEFORE YOU HAVE RUN THE PROBLEM UNDER THE CLOCK. This is Lucene's actual design in miniature, and it should be your fastest problem if retrieval is your background. Three decisions: 1. INTERSECT THE SMALLEST POSTINGS LIST FIRST. Query cost is dominated by the rarest term, because after intersecting with it the candidate set can only shrink. Starting with the most common term does the same work in the worst possible order. 2. DELETION IS A TOMBSTONE PLUS A LIVE-DOCS SET, not removal from every postings list. Removing would be O(terms in the document) with random access across the whole index -- and in a segmented design the older segments are immutable, so it is not even possible. Searches filter against live docs. 3. SEGMENTS ARE IMMUTABLE. Writes accumulate in a mutable buffer; flush seals it. Queries span every segment plus the buffer, newest first, and a delete recorded in a newer segment MASKS a document that an older segment still contains. Merging reclaims the tombstoned space. That is exactly why Lucene can serve queries while indexing. """ from __future__ import annotations import math import re TOKEN_RE = re.compile(r"[a-z0-9]+") K1 = 1.5 B = 0.75 def tokenize(text): return TOKEN_RE.findall(text.lower()) class _Segment: """An immutable slice of the index. `deleted` masks older segments.""" __slots__ = ("postings", "lengths", "deleted") def __init__(self): self.postings = {} # term -> {doc_id: term frequency} self.lengths = {} # doc_id -> token count self.deleted = set() # doc_ids deleted while this segment was live def add(self, doc_id, tokens): self.lengths[doc_id] = len(tokens) counts = {} for token in tokens: counts[token] = counts.get(token, 0) + 1 for term, freq in counts.items(): self.postings.setdefault(term, {})[doc_id] = freq class Index: def __init__(self): self._segments = [] # oldest first self._buffer = _Segment() # the mutable one # ---- writing --------------------------------------------------------- def add(self, doc_id, text): """Re-adding REPLACES. Where the old copy lives decides how.""" owner = self._find_live_segment(doc_id) if owner is self._buffer: # Still in the mutable buffer, so it can genuinely be removed. self._buffer.lengths.pop(doc_id, None) for postings in self._buffer.postings.values(): postings.pop(doc_id, None) self._buffer.deleted.discard(doc_id) elif owner is not None: # It lives in a sealed, immutable segment: mask it instead. self._buffer.deleted.add(doc_id) else: self._buffer.deleted.discard(doc_id) # clear any tombstone self._buffer.add(doc_id, tokenize(text)) def delete(self, doc_id): if not self._is_live(doc_id): return False if doc_id in self._buffer.lengths: # Present in the mutable buffer: it can genuinely be removed. del self._buffer.lengths[doc_id] for postings in self._buffer.postings.values(): postings.pop(doc_id, None) self._buffer.deleted.add(doc_id) return True def flush(self): """Seal the buffer into an immutable segment.""" if self._buffer.lengths or self._buffer.deleted: self._segments.append(self._buffer) self._buffer = _Segment() def merge(self): """Combine every segment, dropping tombstoned documents entirely.""" self.flush() merged = _Segment() for doc_id in self._live_docs(): segment = self._find_live_segment(doc_id) merged.lengths[doc_id] = segment.lengths[doc_id] for term, postings in segment.postings.items(): if doc_id in postings: merged.postings.setdefault(term, {})[doc_id] = postings[doc_id] self._segments = [merged] if merged.lengths else [] self._buffer = _Segment() @property def segment_count(self): return len(self._segments) # ---- liveness -------------------------------------------------------- def _all_segments(self): return self._segments + [self._buffer] def _is_live(self, doc_id): return self._find_live_segment(doc_id) is not None def _find_live_segment(self, doc_id): """The segment holding the authoritative copy, or None if deleted. Walk newest-first: the first segment with an opinion wins, because a newer delete masks an older add. Within one segment an ADD wins over a delete, since re-adding into the buffer supersedes a tombstone it recorded a moment earlier. """ for segment in reversed(self._all_segments()): if doc_id in segment.lengths: return segment if doc_id in segment.deleted: return None return None def _live_docs(self): candidates = set() for segment in self._all_segments(): candidates.update(segment.lengths) return sorted(d for d in candidates if self._is_live(d)) @property def doc_count(self): return len(self._live_docs()) # ---- querying -------------------------------------------------------- def _postings(self, term): """doc_id -> frequency, taken ONLY from each doc's live segment. Unioning across segments and filtering by liveness afterwards is the subtle bug: a document re-added with new content is live, but its OLD segment still holds postings for the OLD terms, so the previous content resurfaces in results. """ out = {} for segment in reversed(self._all_segments()): # newest first for doc_id, freq in segment.postings.get(term, {}).items(): if doc_id in out: continue # a newer segment already won if self._find_live_segment(doc_id) is segment: out[doc_id] = freq return out def search_all(self, query): terms = tokenize(query) if not terms: return [] lists = [set(self._postings(t)) for t in terms] # Smallest first: after intersecting with the rarest term the candidate # set can only shrink, so this is most of the query performance. lists.sort(key=len) result = lists[0] for other in lists[1:]: result &= other if not result: break return sorted(result) def search_any(self, query): result = set() for term in tokenize(query): result |= set(self._postings(term)) return sorted(result) def search(self, query, k=10): terms = tokenize(query) if not terms: return [] live = self._live_docs() n = len(live) if n == 0: return [] lengths = {} for doc_id in live: segment = self._find_live_segment(doc_id) lengths[doc_id] = segment.lengths[doc_id] avgdl = sum(lengths.values()) / n scores = {} for term in terms: postings = self._postings(term) df = len(postings) if df == 0: continue idf = math.log(1 + (n - df + 0.5) / (df + 0.5)) for doc_id, freq in postings.items(): dl = lengths[doc_id] denom = freq + K1 * (1 - B + B * dl / avgdl) scores[doc_id] = scores.get(doc_id, 0.0) + \ idf * (freq * (K1 + 1) / denom) # Ties break by doc_id ascending so results are deterministic. ranked = sorted(scores.items(), key=lambda kv: (-kv[1], kv[0])) return ranked[:k]