74 lines
1.6 KiB
Python
74 lines
1.6 KiB
Python
# coding: utf-8
|
|
from queue import Queue
|
|
|
|
|
|
class Trie:
|
|
""" String trie """
|
|
|
|
def __init__(self):
|
|
self.key = ''
|
|
self.value = None
|
|
self.children = [None] * 26
|
|
self.isEnd = False
|
|
|
|
def insert(self, key: str, value):
|
|
""" insert item """
|
|
key = key.lower()
|
|
|
|
node = self
|
|
for c in key:
|
|
i = ord(c) - 97
|
|
if not 0 <= i < 26:
|
|
return
|
|
|
|
if not node.children[i]:
|
|
node.children[i] = Trie()
|
|
|
|
node = node.children[i]
|
|
|
|
node.isEnd = True
|
|
node.key = key
|
|
node.value = value
|
|
|
|
def get(self, key, default=None):
|
|
""" get value of key """
|
|
node = self.searchPrefix(key)
|
|
if not (node and node.isEnd):
|
|
return default
|
|
|
|
return node.value
|
|
|
|
def searchPrefix(self, prefix):
|
|
""" search node matchs the prefix """
|
|
prefix = prefix.lower()
|
|
node = self
|
|
for c in prefix:
|
|
i = ord(c) - 97
|
|
if not (0 <= i < 26 and node.children[i]):
|
|
return None
|
|
|
|
node = node.children[i]
|
|
|
|
return node
|
|
|
|
def items(self, prefix):
|
|
""" search items match the prefix """
|
|
node = self.searchPrefix(prefix)
|
|
if not node:
|
|
return []
|
|
|
|
q = Queue()
|
|
result = []
|
|
q.put(node)
|
|
|
|
while not q.empty():
|
|
node = q.get()
|
|
if node.isEnd:
|
|
result.append((node.key, node.value))
|
|
|
|
for c in node.children:
|
|
if c:
|
|
q.put(c)
|
|
|
|
return result
|