trjhtr

This is just a classic implementation of a trie using python 3.6, with types. I am a PHP developer, so these are my first attempts with Python.

I know I can improve the readability of my code, it's a bit messy, but I am struggling finding the right way of doing it. I'm sure I am missing a lot of the Pythonic ways in here :-)

from typing import Optional


class Trie:
 def __init__(self):
 self.root = Node(char=None, is_word=False)

 def add(self, word: str) -> None:
 current = self.root
 is_end_of_word = False

 for i in range(0, len(word)):
 char = word[i]
 if i == len(word) - 1:
 is_end_of_word = True

 if not current.children:
 node_to_insert = Node(char, is_end_of_word)
 current.add_child(node_to_insert)
 current = node_to_insert
 continue

 if char not in current.children:
 node_to_insert = Node(char, is_end_of_word)
 current.add_child(node_to_insert)
 current = node_to_insert
 else:
 current = current.children[char]

 def contains(self, word: str) -> bool:
 current = self.root

 for char in word:
 if not current.children:
 return False

 if char in current.children:
 current = current.children[char]
 continue
 else:
 return False

 if not current.is_word:
 return False

 return True

 def remove(self, word: str) -> None:
 current = self.root

 for i in range(0, len(word)):
 char = word[i]
 is_end_of_word = False

 if i == len(word) - 1:
 is_end_of_word = True

 if char in current.children:
 if current.children[char].is_word and is_end_of_word:
 current.children[char].is_word = False
 return
 current = current.children[char]
 else:
 return

 def retrieve_all_words(self) -> list:
 return self._retrieve_all_words(self.root, '', )

 def _retrieve_all_words(self, current: 'Node', word: str, words: list) -> list:
 for child in current.children.values():
 word = word + child.char
 if child.is_word:
 words.append(word)

 if child.children:
 self._retrieve_all_words(child, word, words)
 word = word[:-1]
 continue

 word = word[:-1]

 return words


class Node:
 def __init__(self, char: Optional[str], is_word: bool):
 self.char = char
 self.is_word = is_word
 self.children = 

 def add_child(self, node: 'Node'):
 self.children[node.char] = node

Your code looks very good and the type annotation is a really nice touch
Also, your code is well organised and easy to understand. Some documentation could be a good idea (but it's better to have no doc rather than bad doc). Also, you could consider writing unit tests.

Let's try to see what could be improved/made more Pythonic.

Loop like a native

I highly recommand Ned Batchelder's excellent talk "loop like a native". One of the idea is that whenever you are using for i in range(len(iterable)), you are probably doing it wrong. In your case, you could use for i, char in enumerate(word):.

End of word

The end-of-word detection could be done in a single statement: is_end_of_word = i == len(word) - 1. Also, you can get rid of the definition before the loop and even in the loops, sometimes, you could define it only behind the if char in current.children: because you use it only there.

Reorganise the logic

Sometimes, you check if something is empty and then if it contains a particular element. This can be factorised out:

Also, simplifying the code if (cond) return False else return True, you could rewrite contains:

def contains(self, word: str) -> bool:
 current = self.root
 for char in word:
 if char not in current.children:
 return False
 current = current.children[char]
 return current.is_word

Then in _retrieve_all_words, you can get rid of continue by using a simple else which makes things more explicit. Then, you can actually factorise out the common code at the end of the 2 branches and get the more simple:

 if child.is_word:
 words.append(word)
 if child.children:
 self._retrieve_all_words(child, word, words)
 word = word[:-1]

Finally, you can use += to simplify word = word + child.char into word += child.char.

At this stage, the code looks like:

from typing import Optional

class Trie:
 def __init__(self):
 self.root = Node(char=None, is_word=False)

 def add(self, word: str) -> None:
 current = self.root

 for i, char in enumerate(word):
 if char in current.children:
 current = current.children[char]
 else:
 is_end_of_word = i == len(word) - 1 
 node_to_insert = Node(char, is_end_of_word)
 current.add_child(node_to_insert)
 current = node_to_insert

 def contains(self, word: str) -> bool:
 current = self.root
 for char in word:
 if char not in current.children:
 return False
 current = current.children[char]
 return current.is_word

 def remove(self, word: str) -> None:
 current = self.root
 for i, char in enumerate(word):
 if char not in current.children:
 return
 is_end_of_word = i == len(word) - 1 
 if current.children[char].is_word and is_end_of_word:
 current.children[char].is_word = False
 return
 current = current.children[char]

 def retrieve_all_words(self) -> list:
 return self._retrieve_all_words(self.root, '', )

 def _retrieve_all_words(self, current: 'Node', word: str, words: list) -> list:
 for child in current.children.values():
 word += child.char
 if child.is_word:
 words.append(word)
 if child.children:
 self._retrieve_all_words(child, word, words)
 word = word[:-1]
 return words


class Node:
 def __init__(self, char: Optional[str], is_word: bool):
 self.char = char
 self.is_word = is_word
 self.children = 

 def add_child(self, node: 'Node'):
 self.children[node.char] = node

t = Trie()
t.add("toto")
t.add("tutu")
t.add("foobar")
print(t)
print(t.retrieve_all_words())
print(t.contains("tot"))
print(t.contains("toto"))
print(t.contains("totot"))

Storing things in a different format

Instead of maintaining a is_word attribute, maybe you could add store a sentinel like None to say that we have a full word here. This could be written like this:

from typing import Optional

class Trie:
 def __init__(self):
 self.root = Node(char=None)

 def add(self, word: str) -> None:
 current = self.root

 for char in list(word) + [None]:
 if char in current.children:
 current = current.children[char]
 else:
 node_to_insert = Node(char)
 current.add_child(node_to_insert)
 current = node_to_insert

 def contains(self, word: str) -> bool:
 current = self.root
 for char in list(word) + [None]:
 if char not in current.children:
 return False
 current = current.children[char]
 return True 

 def remove(self, word: str) -> None:
 current = self.root
 for char in list(word) + [None]:
 if char not in current.children:
 return
 elif char is None:
 del current.children[char]
 else:
 current = current.children[char]

 def retrieve_all_words(self) -> list:
 return self._retrieve_all_words(self.root, '', )

 def _retrieve_all_words(self, current: 'Node', word: str, words: list) -> list:
 for child in current.children.values():
 if child.char is None:
 words.append(word)
 else:
 word += child.char
 if child.children:
 self._retrieve_all_words(child, word, words)
 word = word[:-1]
 return words


class Node:
 def __init__(self, char: Optional[str]):
 self.char = char
 self.children = 

 def add_child(self, node: 'Node'):
 self.children[node.char] = node

t = Trie()
t.add("toto")
t.add("tutu")
t.add("foobar")
print(t)
print(t.retrieve_all_words())
print(t.contains("tot"))
print(t.contains("toto"))
print(t.contains("totot"))
t.remove("toto")
print(t.retrieve_all_words())
print(t.contains("toto"))

Back to word = word + child.char

Something I should have mentionned earlier but spotted on the late.

In _retrieve_all_words, you add a char to word only to remove it afterward. It is much clearer (and more efficient?) to just write word + char in the places where you actually need to word with the added character.

Then, the code becomes:

def _retrieve_all_words(self, current: 'Node', word: str, words: list) -> list:
 for child in current.children.values():
 if child.char is None:
 words.append(word)
 elif child.children:
 self._retrieve_all_words(child, word + child.char, words)
 return words

this is a small thing, but the API you provide isn't very pythonic. rather than using a containsmethod, you should overload in. similarly, rather than defining return_all_words, you should define iteration, so you can loop directly through it, and then the list conversion will just be list(tried). these may seem insignificant, but this type of consistency is what makes python feel nice to use.