class Trie(object):
def __init__(self):
self.root = TrieNode()
def insert(self, data):
node = self.root
idx = 0
for char in data:
if char not in node.children:
node.children[char] = TrieNode()
node = node.children[char]
if len(node.search_words) < 3:
node.search_words.append(data)
idx += 1
def search(self, word):
result, node = [], self.root
for i, char in enumerate(word):
if char not in node.children:
temp = [[] for _ in range(len(word) - i)]
return result + temp
else:
node = node.children[char]
result.append(node.search_words[:])
return result