trie.py 597 B

1234567891011121314151617181920212223242526272829
  1. # Copyright (c) Alibaba, Inc. and its affiliates.
  2. from collections import defaultdict
  3. class TreeNode:
  4. def __init__(self):
  5. self.child = defaultdict(TreeNode)
  6. class Trie:
  7. def __init__(self, eos):
  8. self.root = TreeNode()
  9. self.eos = eos
  10. def insert(self, word):
  11. cur = self.root
  12. for c in word:
  13. cur = cur.child[c]
  14. def get_next_layer(self, word):
  15. cur = self.root
  16. for c in word:
  17. cur = cur.child.get(c)
  18. if cur is None:
  19. return [self.eos]
  20. return list(cur.child.keys())