troupe-standup/wizard-work/beam/test_trie.py

"""Cells for trie.Trie.

Every cell is a unittest.TestCase method so `python3 -m unittest discover`
collects it. The module imported is pinned to the file beside this one: an
environment that shadows `trie` (a host copy, an installed package) raises
instead of grading the wrong file green -- Q68's fault, measured in this lane.

`python3 test_trie.py` still runs the same cells from __main__.
"""

import os
import random
import sys
import unittest

_HERE = os.path.dirname(os.path.abspath(__file__))
if sys.path and sys.path[0] != _HERE:
    sys.path.insert(0, _HERE)

import trie as _trie_module
from trie import Trie

# abspath, not realpath: realpath would resolve a symlinked module to the same
# string as its target and let a symlink-out pass as "the file beside me".
_EXPECTED = os.path.abspath(os.path.join(_HERE, "trie.py"))
_GOT = os.path.abspath(_trie_module.__file__)
if _GOT != _EXPECTED:
    raise ImportError(
        "trie resolved to %r, expected %r -- these cells would grade the wrong file"
        % (_GOT, _EXPECTED)
    )


class TestBasic(unittest.TestCase):
    def test_basic(self):
        t = Trie()
        for w in ["cat", "car", "card", "dog", ""]:
            t.insert(w)
        self.assertTrue(t.search("cat"))
        self.assertTrue(t.search(""))
        self.assertTrue(t.search("card"))
        self.assertFalse(t.search("ca"))
        self.assertTrue(t.starts_with("ca"))
        self.assertFalse(t.starts_with("z"))
        self.assertEqual(t.words_with_prefix("ca"), ["car", "card", "cat"])
        self.assertEqual(t.count_prefix("ca"), 3)
        self.assertEqual(len(t), 5)

    def test_module_resolved_beside_these_cells(self):
        self.assertEqual(_GOT, _EXPECTED)


class TestDelete(unittest.TestCase):
    def test_delete(self):
        t = Trie()
        for w in ["car", "card", "care"]:
            t.insert(w)
        self.assertTrue(t.delete("car"))
        self.assertFalse(t.search("car"))
        self.assertTrue(t.search("card"))
        self.assertTrue(t.search("care"))
        self.assertTrue(t.starts_with("car"))
        self.assertFalse(t.delete("car"))
        # deleting the shared-prefix leaves should prune the branch
        t.delete("card")
        t.delete("care")
        self.assertFalse(t.starts_with("car"))

    def test_delete_prunes_counters_with_it(self):
        t = Trie()
        for w in ["car", "card", "care"]:
            t.insert(w)
        t.delete("card")
        t.delete("care")
        self.assertEqual(t.count_prefix("car"), 1)  # only "car" left below it
        self.assertEqual(t.count_prefix("car"), len(t.words_with_prefix("car")))
        t.delete("car")
        self.assertEqual(t.count_prefix(""), 0)  # branch pruned, root empty
        self.assertEqual(len(t), 0)


class TestDuplicateInsert(unittest.TestCase):
    def test_duplicate_insert(self):
        t = Trie()
        t.insert("abc")
        t.insert("abc")
        self.assertTrue(t.search("abc"))
        t.delete("abc")
        # one delete removes the word entirely: multiplicity is not tracked,
        # so the counter must read 1, not 2, for the whole lifetime
        self.assertFalse(t.search("abc"))
        self.assertEqual(len(t), 0)
        self.assertEqual(t.count_prefix(""), 0)

    def test_duplicate_insert_does_not_inflate_counts(self):
        t = Trie()
        t.insert("abc")
        t.insert("abc")
        self.assertEqual(len(t), 1)
        self.assertEqual(t.count_prefix("a"), 1)
        self.assertEqual(t.count_prefix("ab"), 1)
        self.assertEqual(t.count_prefix("abc"), 1)
        self.assertEqual(t.count_prefix("a"), len(t.words_with_prefix("a")))


class TestCountPrefixAgreesWithEnumeration(unittest.TestCase):
    """The item's cell: the counter may not disagree with the enumeration."""

    def assert_agrees(self, t, prefixes):
        for p in prefixes:
            self.assertEqual(
                t.count_prefix(p),
                len(t.words_with_prefix(p)),
                "prefix %r: counter %d vs enumerated %d"
                % (p, t.count_prefix(p), len(t.words_with_prefix(p))),
            )

    def test_after_delete_of_a_word(self):
        t = Trie()
        for w in ["cat", "car", "card", "dog", ""]:
            t.insert(w)
        t.delete("car")
        self.assert_agrees(t, ["", "c", "ca", "car", "cat", "card", "d", "do", "dog", "z"])

    def test_after_delete_of_a_prefix_word(self):
        t = Trie()
        for w in ["a", "ab", "abc"]:
            t.insert(w)
        t.delete("a")
        self.assertEqual(len(t), 2)
        self.assert_agrees(t, ["", "a", "ab", "abc"])

    def test_after_insert_then_delete_of_everything(self):
        t = Trie()
        t.insert("zzz")
        t.delete("zzz")
        self.assertEqual(len(t), 0)
        self.assert_agrees(t, ["", "z", "zz", "zzz"])

    def test_root_count_equals_len(self):
        t = Trie()
        for w in ["a", "ab", "abc", "", "b"]:
            t.insert(w)
        self.assertEqual(t.count_prefix(""), len(t))
        self.assertEqual(t.count_prefix(""), 5)

    def test_interleaved_insert_delete_random(self):
        rng = random.Random(1234)
        t = Trie()
        live = set()
        prefixes = ["", "a", "b", "c", "aa", "ab", "ba", "cc"]
        for _ in range(400):
            w = "".join(rng.choice("abc") for _ in range(rng.randint(0, 3)))
            if rng.random() < 0.6:
                t.insert(w)
                live.add(w)
            else:
                removed = t.delete(w)
                self.assertEqual(removed, w in live, "delete(%r) returned %r" % (w, removed))
                live.discard(w)
            self.assertEqual(len(t), len(live))
            self.assert_agrees(t, prefixes)
            for p in prefixes:
                self.assertEqual(
                    set(t.words_with_prefix(p)), {x for x in live if x.startswith(p)}
                )


class TestRandomVsBruteforce(unittest.TestCase):
    def test_random_vs_bruteforce(self):
        rng = random.Random(123)
        words = []
        alphabet = "ab"
        for _ in range(200):
            w = "".join(rng.choice(alphabet) for _ in range(rng.randint(0, 5)))
            words.append(w)
        t = Trie()
        s = set()
        for w in words:
            t.insert(w)
            s.add(w)
        for w in set(words):
            self.assertEqual(t.search(w), w in s)
        # delete half
        to_del = list(s)[: len(s) // 2]
        for w in to_del:
            self.assertTrue(t.delete(w))
            s.discard(w)
        for w in set(words):
            self.assertEqual(t.search(w), w in s, (w, t.search(w), w in s))
        # prefix enumeration matches
        for p in ["", "a", "b", "aa", "ab"]:
            self.assertEqual(set(t.words_with_prefix(p)), {w for w in s if w.startswith(p)})
        # and the counter agrees with it after the deletes
        for p in ["", "a", "b", "aa", "ab"]:
            self.assertEqual(t.count_prefix(p), len(t.words_with_prefix(p)))
        self.assertEqual(t.count_prefix(""), len(t))


if __name__ == "__main__":
    unittest.main(verbosity=2)