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") == [])