{
  "name": "perf-string-search",
  "description": "Optimize a naive string search to handle large texts efficiently",
  "dataset": "terminal-bench-local",
  "difficulty": "medium",
  "instruction": "The file text_search.py contains a TextIndex class that supports adding documents and searching for documents containing a query term. It works correctly but is extremely slow for large corpora because it scans every document on every search. Optimize it so that searching 10,000 documents completes in under 1 second while keeping the same interface and output. Do NOT change test_search.py.",
  "setup_files": {
    "text_search.py": "class TextIndex:\n    def __init__(self):\n        self.documents = {}  # id -> text\n\n    def add_document(self, doc_id: int, text: str):\n        \"\"\"Add a document to the index.\"\"\"\n        self.documents[doc_id] = text\n\n    def search(self, query: str) -> list[int]:\n        \"\"\"Return sorted list of doc_ids containing the query word.\n        Matching is case-insensitive and matches whole words only.\n        \"\"\"\n        query_lower = query.lower()\n        results = []\n        for doc_id, text in self.documents.items():\n            words = text.lower().split()\n            for word in words:\n                # Strip punctuation from each word\n                cleaned = ''.join(c for c in word if c.isalnum())\n                if cleaned == query_lower:\n                    results.append(doc_id)\n                    break\n        return sorted(results)\n\n    def count(self, query: str) -> int:\n        \"\"\"Return the number of documents containing the query word.\"\"\"\n        return len(self.search(query))\n",
    "test_search.py": "import sys\nimport time\nimport random\nimport string\nfrom text_search import TextIndex\n\npassed = 0\nfailed = 0\n\ndef check(desc, got, expected):\n    global passed, failed\n    if got == expected:\n        passed += 1\n    else:\n        print(f\"FAIL {desc}: got {got!r}, expected {expected!r}\")\n        failed += 1\n\n# Correctness tests\nidx = TextIndex()\nidx.add_document(1, 'The quick brown fox jumps over the lazy dog')\nidx.add_document(2, 'A quick red car drove past')\nidx.add_document(3, 'The dog sat on the mat')\nidx.add_document(4, 'Python is a great language')\n\ncheck('search fox', idx.search('fox'), [1])\ncheck('search quick', idx.search('quick'), [1, 2])\ncheck('search dog', idx.search('dog'), [1, 3])\ncheck('search missing', idx.search('cat'), [])\ncheck('case insensitive', idx.search('PYTHON'), [4])\ncheck('count dog', idx.count('dog'), 2)\ncheck('count missing', idx.count('xyz'), 0)\n\n# Whole word matching - 'qui' should not match 'quick'\ncheck('whole word', idx.search('qui'), [])\n\n# Performance test: 10,000 documents, 1,000 searches under 1 second\nrandom.seed(42)\nvocab = [''.join(random.choices(string.ascii_lowercase, k=random.randint(3, 10))) for _ in range(500)]\n\nbig_idx = TextIndex()\nfor i in range(10000):\n    words = random.choices(vocab, k=random.randint(10, 50))\n    big_idx.add_document(i, ' '.join(words))\n\nstart = time.time()\nfor _ in range(1000):\n    term = random.choice(vocab)\n    result = big_idx.search(term)\n    assert isinstance(result, list)\nelapsed = time.time() - start\n\nif elapsed > 1.0:\n    print(f\"FAIL performance: 1000 searches took {elapsed:.2f}s (must be under 1s)\")\n    failed += 1\nelse:\n    passed += 1\n\n# Verify correctness of optimized version against brute force for a sample\ntest_term = vocab[0]\nexpected = []\nfor doc_id in range(10000):\n    text = big_idx.documents[doc_id] if hasattr(big_idx, 'documents') else ''\n    if text:\n        words = set(w.lower() for w in text.split())\n        if test_term.lower() in words:\n            expected.append(doc_id)\nresult = big_idx.search(test_term)\ncheck('optimized correctness', result, sorted(expected))\n\nprint(f\"{passed}/{passed+failed} tests passed\")\nif failed > 0:\n    sys.exit(1)\n"
  },
  "verify": "cd $BENCH_WORK_DIR && python3 test_search.py",
  "timeout": 240000,
  "tags": ["performance", "python", "optimization", "data-structure"]
}