74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
import random
|
|
from trie import Trie
|
|
|
|
|
|
def test_basic():
|
|
t = Trie()
|
|
for w in ["cat", "car", "card", "dog", ""]:
|
|
t.insert(w)
|
|
assert t.search("cat")
|
|
assert t.search("")
|
|
assert t.search("card")
|
|
assert not t.search("ca")
|
|
assert t.starts_with("ca")
|
|
assert not t.starts_with("z")
|
|
assert sorted(t.words_with_prefix("ca")) == ["car", "card", "cat"]
|
|
assert t.count_prefix("ca") == 3
|
|
assert len(t) == 5
|
|
|
|
|
|
def test_delete():
|
|
t = Trie()
|
|
for w in ["car", "card", "care"]:
|
|
t.insert(w)
|
|
assert t.delete("car")
|
|
assert not t.search("car")
|
|
assert t.search("card")
|
|
assert t.search("care")
|
|
assert t.starts_with("car")
|
|
assert not t.delete("car")
|
|
# delete shared prefix leaf should prune
|
|
t.delete("card"); t.delete("care")
|
|
assert not t.starts_with("car")
|
|
|
|
|
|
def test_duplicate_insert():
|
|
t = Trie()
|
|
t.insert("abc"); t.insert("abc")
|
|
assert t.search("abc")
|
|
t.delete("abc")
|
|
assert not t.search("abc") # one delete removes word entirely (count not tracked per multiplicity)
|
|
# acceptable: delete removes end flag
|
|
|
|
|
|
def test_random_vs_bruteforce():
|
|
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):
|
|
assert t.search(w) == (w in s)
|
|
# delete half
|
|
to_del = list(s)[: len(s) // 2]
|
|
for w in to_del:
|
|
assert t.delete(w)
|
|
s.discard(w)
|
|
for w in set(words):
|
|
assert t.search(w) == (w in s), (w, t.search(w), w in s)
|
|
# prefix enumeration matches
|
|
for p in ["", "a", "b", "aa", "ab"]:
|
|
got = set(t.words_with_prefix(p))
|
|
exp = {w for w in s if w.startswith(p)}
|
|
assert got == exp, (p, got, exp)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_basic(); test_duplicate_insert(); test_delete(); test_random_vs_bruteforce()
|
|
print("ALL OK")
|