def test_removePrefix():
trie = Trie()
trie.add("abcd")
trie.add("abc")
trie.add("abd")
trie.add("abce")
trie.add("words")
trie.add("swords")
assert (trie.count("abc") == 3)
assert (trie.count("a") == 4)
assert (len(trie) == 6)
trie.removePrefix("abc") # removes 'abc', 'abcd' and 'abce'
# 'ab' should only have 'd' as its child now
assert (len(trie) == 3)
assert (trie.count("abc") == 0)
assert (trie.count("a") == 1)
assert (trie.findWords("a") == ["abd"])
# try and remove 'abc' again - nothing should change
trie.removePrefix("abc") # NO-OP
assert (len(trie) == 3)
assert (trie.count("abc") == 0)
assert (trie.count("a") == 1)
assert (trie.findWords("a") == ["abd"])
# remove 'a*'
trie.removePrefix("a") # removes 'abd'
assert (len(trie) == 2)
assert (trie.count("abc") == 0)
assert (trie.count("a") == 0)
assert (trie.findWords("a") == [])
assert (trie.findWords("") == ["swords", "words"])
# flush the whole trie
trie.removePrefix("")
assert (trie.root is not None)
assert (len(trie) == 0)
trie.add("words")
trie.add("newly")
trie.add("added")
assert (len(trie) == 3)
# flush the whole trie again
trie.removePrefix("")
assert (trie.root is not None)
assert (len(trie) == 0)
def test_findWords():
trie = Trie()
trie.add("cdef")
trie.add("abcd")
trie.add("abc")
trie.add("acdc")
trie.add("bcdef")
assert (len(trie) == 5)
assert (trie.findWords("abc") == ["abc", "abcd"])
assert (trie.findWords("a") == ["abc", "abcd", "acdc"])
assert (trie.findWords("") == ["abc", "abcd", "acdc", "bcdef", "cdef"])
assert (trie.findWords() == ["abc", "abcd", "acdc", "bcdef", "cdef"])
assert (trie.findWords("ad") == [])
assert (trie.findWords("b") == ["bcdef"])
assert (trie.findWords("c") == ["cdef"])
assert (trie.findWords("d") == [])