Trie implementation with Python 3 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) -&gt; 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) -&gt; 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) -&gt; 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) -&gt; list:
        return self._retrieve_all_words(self.root, '', [])

    def _retrieve_all_words(self, current: 'Node', word: str, words: list) -&gt; 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) -&gt; 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) -&gt; 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) -&gt; 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) -&gt; 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) -&gt; list:
        return self._retrieve_all_words(self.root, '', [])

    def _retrieve_all_words(self, current: 'Node', word: str, words: list) -&gt; 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) -&gt; 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) -&gt; 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) -&gt; 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) -&gt; list:
        return self._retrieve_all_words(self.root, '', [])

    def _retrieve_all_words(self, current: 'Node', word: str, words: list) -&gt; 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) -&gt; 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