Guest User

Untitled

a guest
Mar 26th, 2018
101
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. from typing import Optional
  2.  
  3. class Trie:
  4.     def __init__(self):
  5.         self.root = Node(None)
  6.  
  7.     def add(self, word: str) -> None:
  8.         current = self.root
  9.  
  10.         for i, char in enumerate(word):
  11.             if char in current.children:
  12.                 current = current.children[char]
  13.             else:
  14.                 node_to_insert = Node(char)
  15.                 current.add_child(node_to_insert)
  16.                 current = node_to_insert
  17.         current.is_word = True
  18.  
  19.     def contains(self, word: str) -> bool:
  20.         current = self.root
  21.         for char in word:
  22.             if char not in current.children:
  23.                 return False
  24.             current = current.children[char]
  25.         return current.is_word
  26.  
  27.     def remove(self, word: str) -> None:
  28.         current = self.root
  29.         for i, char in enumerate(word):
  30.             if char not in current.children:
  31.                 return
  32.             is_end_of_word = i == len(word) - 1
  33.             if current.children[char].is_word and is_end_of_word:
  34.                 current.children[char].is_word = False
  35.                 return
  36.             current = current.children[char]
  37.  
  38.     def retrieve_all_words(self) -> list:
  39.         return self._retrieve_all_words(self.root, '', [])
  40.  
  41.     def _retrieve_all_words(self, current: 'Node', word: str, words: list) -> list:
  42.         for child in current.children.values():
  43.             word += child.char
  44.             if child.is_word:
  45.                 words.append(word)
  46.             if child.children:
  47.                 self._retrieve_all_words(child, word, words)
  48.             word = word[:-1]
  49.         return words
  50.  
  51.  
  52. class Node:
  53.     def __init__(self, char: Optional[str]):
  54.         self.char = char
  55.         self.is_word = False
  56.         self.children = {}
  57.  
  58.     def add_child(self, node: 'Node'):
  59.         self.children[node.char] = node
  60.  
  61. t = Trie()
  62. t.add("toto")
  63. t.add("tutu")
  64. t.add("foobar")
  65. print(type(t) == Trie)
  66. print(len(t.retrieve_all_words()) == 3)
  67. print(t.contains("tot") == False)
  68. print(t.contains("toto") == True)
  69. print(t.contains("totot") == False)
RAW Paste Data