Trie implementation in Python

Question

I have created a trie in Python and am looking for feedback. Specifically, I am looking for feedback on:

1. If my code is 'pythonic'
2. If my logic is proper when inserting and retrieving from the trie

I have implemented the following logic:

1. Insert a word into the trie
2. Check if a word exists in the trie
3. Given a prefix, get all possible words . Get all the words in the trie

My purpose for coding the trie is to improve my Python and algorithm knowledge.

from collections import deque


class Node:
    def __init__(self, character, parent):
        self.character = character
        if self.character is not None:
            self.character = self.character.lower()
        self.parent = parent
        self.children = dict()
        self.terminus = False

    def add(self, child_node):
        self.children[child_node.character] = child_node


class Trie:

    def __init__(self):
        self._root = Node(None, None)

    def insert(self, word):
        if word:
            current_node = self._root
            for i, character in enumerate(self._normalize_word(word)):
                if character in current_node.children:
                    current_node = current_node.children[character]
                else:
                    child_node = Node(character, current_node)
                    current_node.add(child_node)
                    current_node = child_node
            current_node.terminus = True

    def __contains__(self, item):
        current_node = self._root
        contained = True
        for symbol in self._normalize_word(item):
            if symbol in current_node.children:
                current_node = current_node.children[symbol]
            else:
                contained = False
                break
        return contained and current_node.terminus

    def _normalize_word(self, word):
        return word.strip().lower()

    def _get_all_words(self, prefix, node, word_list):
        if node.character:
            prefix.append(node.character)
        for child in node.children.values():
            self._get_all_words(prefix, child, word_list)
        if node.terminus:
            word_list.append("".join([i[0] for i in prefix]))
        if len(prefix) > 0:
            prefix.pop()

    def get_possible_words(self, prefix):
        current_node = self._root
        found_prefix = True
        word_list = []
        prefix_deque = deque()
        for symbol in prefix.strip().lower():
            if symbol in current_node.children:
                current_node = current_node.children[symbol]
            else:
                found_prefix = False
                break
        if found_prefix:
            self._get_all_words(prefix_deque, current_node, word_list)
            # This is a bit kludgy - add the prefix to the rest of the characters found,
            # but I take off the last character from the prefix because it was added
            # in the _get_all_words method call since it is the current_node being passed into it. 
            # Is there a better way to do this?
        word_list = list(map(lambda word: prefix[:len(prefix)-1] + word, word_list))

        return word_list

    def get_all_words(self):
        word_list = []
        self._get_all_words(deque(), self._root, word_list)
        return word_list

Here is my unit test:

import unittest
from algorithms import trie


class TestTrie(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        #            root
        #          /      \
        #       /           \
        #     a *             b
        #    / \             /
        #   d   n *         a
        #  /   / \       /  |  \
        # d * d * y *   g * t *  y *
        #              / \   \
        #             e   s * h *
        #            /
        #           l *
        # asterisk denotes a word

        cls._trie = trie.Trie()
        cls._trie.insert("a")
        cls._trie.insert("add")
        cls._trie.insert("an")
        cls._trie.insert("and")
        cls._trie.insert("any")
        cls._trie.insert("bagel")
        cls._trie.insert("bag")
        cls._trie.insert("bags")
        cls._trie.insert("bat")
        cls._trie.insert("bath")
        cls._trie.insert("bay")

        cls._trie_length = 11 # magic number, the number of words in the trie

    def test(self):
        assert len(self._trie.get_all_words()) == self._trie_length

        assert "a" in self._trie
        assert "add" in self._trie
        assert "an" in self._trie
        assert "and" in self._trie
        assert "any" in self._trie
        assert "bagel" in self._trie
        assert "bag" in self._trie
        assert "bags" in self._trie
        assert "bat" in self._trie
        assert "bath" in self._trie
        assert "bay" in self._trie

    def test_duplicate_entries(self):
        """Adding a word that already exists should not create a new word in the trie"""
        t = self._trie
        t.insert("bag")

        assert len(t.get_all_words()) == self._trie_length
        assert "bag" in t

    def test_mixed_case(self):
        """insert and retrieval are case insensitive"""
        t = trie.Trie()
        t.insert("APPLE")
        t.insert("oRANge")

        assert "apple" in t
        assert "orange" in t

        assert "APPLE" in t
        assert "ORANGE" in t

        assert "aPpLe" in t
        assert "oRangE" in t

    def test_hyphenated_words(self):
        t = trie.Trie()
        t.insert("e-mail")
        t.insert("above-said")
        t.insert("above-water")
        t.insert("above-written")
        t.insert("above")
        t.insert("abode")
        t.insert("exit")

        assert len(t.get_all_words()) == 7

        assert "abode" in t
        assert "above" in t
        assert "above-written" in t
        assert "above-water" in t
        assert "above-said" in t
        assert "e-mail" in t
        assert "exit" in t

    def test_empty_trie(self):
        t = trie.Trie()
        assert len(t.get_all_words()) == 0

    def test_first_symbol_is_a_word(self):
        t = trie.Trie()
        t.insert("a")
        t.insert("apple")

        assert "a" in t
        assert "apple" in t

        words = t.get_all_words()
        assert len(words) == 2
        assert "a" in words
        assert "apple" in words

    def test_get_possible_words(self):
        prefix = 'an'
        expected_words = ['an', 'and', 'any']
        actual_words = self._trie.get_possible_words(prefix)
        assert len(expected_words) == len(actual_words)
        for word in expected_words:
            assert word in actual_words

        prefix = 'ba'
        expected_words = ["bagel", "bag", "bags", "bat", "bath", "bay"]
        actual_words = self._trie.get_possible_words(prefix)
        assert len(expected_words) == len(actual_words)
        for word in expected_words:
            assert word in actual_words

    def test_get_possible_words_no_more_words(self):
        """test that given a prefix that is a terminus with no children in the trie, returns that one word"""
        prefix = 'any'
        actual_words = self._trie.get_possible_words(prefix)
        assert len(actual_words) == 1
        assert prefix in actual_words

    def test_get_possible_words_prefix_not_in_trie(self):
        prefix = 'z'
        actual_words = self._trie.get_possible_words(prefix)
        assert len(actual_words) == 0

Answer 1

Starting in the Node class, there was a parent attribute that never gets referenced. Also I find if you are just assigning a variable based off one condition doing it in one line makes the code cleaner and easier to read.

class Node:

    def __init__(self, character):
        self.character = character if character is None else character.lower()
        self.children = dict()
        self.terminus = False

    def add(self, child_node):
        self.children[child_node.character] = child_node

In your Trie class specifically with the .insert and __contains__ methods:

def insert(self, word):
    if word: # checking for truth to continue
        current_node = self._root
        for i, character in enumerate(self._normalize_word(word)):
            if character in current_node.children:
                current_node = current_node.children[character]
            else:
                child_node = Node(character, current_node)
                current_node.add(child_node)
                current_node = child_node
        current_node.terminus = True

You can save a level of indentation by returning on false, essentially short-circuiting your method:

def insert(self, word):
    if not word:
        return
    current_node = self._root
    for i, character in enumerate(self._normalize_word(word)):
        if character in current_node.children:
            current_node = current_node.children[character]
        else:
            child_node = Node(character, current_node)
            current_node.add(child_node)
            current_node = child_node
    current_node.terminus = True

You also make a call to enumerate but never use it:

def insert(self, word):
    if not word:
        return
    current_node = self._root
    for character in self._normalize_word(word):
        if character in current_node.children:
            current_node = current_node.children[character]
        else:
            child_node = Node(character, current_node)
            current_node.add(child_node)
            current_node = child_node
    current_node.terminus = True

Seeing your logic in your method made me think you could also have used a recursive method, something along the lines of:

def insert(self, word):
    if not word: 
        return
    self.__insert(self._normalize_word(word), self._root).terminus = True

def __insert(self, word, node):
    if not word:
        return node
    elif {word[0]} & node.children.keys():
        return self.__insert(word[1:], node.children[word[0]])
    else:
        new_node = Node(word[0])
        node.add(new_node)
        return self.__insert(word[1:], new_node)

One other change I made which depending on the size of the tree will improve performance is the use of a logical & on the children.keys(). PEP 3106 brought about changes that allowed for set operations on dict-keys. This means you could just as well check if a dict has a value as a key by {value} & dict.keys() rather than iterating over every key.

In the __contains__ method where you set a truth flag to false then break and then test with a logical &:

def __contains__(self, item):
    current_node = self._root
    contained = True
    for symbol in self._normalize_word(item):
        if symbol in current_node.children:
            current_node = current_node.children[symbol]
        else:
            contained = False
            break
    return contained and current_node.terminus

If contained is set to False theres no need to check at the bottom of the method. Just get rid of contained and return False if your first condition isn't met. Similar to the last point in .insert you can save a couple of lines if you short-circuit, basically no need for else if you are returning on it:

def __contains__(self, item):
    current_node = self._root
    for symbol in self._normalize_word(item):
        if symbol not in current_node.children:
            return False
        current_node = current_node.children[symbol]
    return current_node.terminus

I saw this as another place where you could have done this recursively as well:

def __contains(self, word, node):
    if not word:
        return node
    return self.__contains(word[1:], node.children[word[0]]) if \
           {word[0]} & node.children.keys() else None

def __contains__(self, item):
    word = self._normalize_word(item)
    return self.__contains(word, self._root) is not None

In the __contains method you can return the Node of the last letter or None. This is useful where you can get a boolean value from your dunder-contains method to check for membership but also you can use the __contains method with the get_possible_words method later. This is similar to how I adjusted the __insert method to return the final Node after inserting to set .terminus.

def get_possible_words(self, word):
    word_node = self.__contains(word, self._root)
    return [] if word_node is None else \
           self.__get_possible_words(word, word_node)

Here you can get your final Node and that will be the starting point to retrieve all possible words. If its None just return an empty list else:

def __get_possible_words(self, word, word_node, word_list=None):
    word_list = [] if word_list is None else word_list
    if word_node.terminus:
        word_list.append(word)
    for letter in word_node.children:
        if not word_node.children[letter]:
            word_list.append(word + letter)
        else:
            self.__get_possible_words(word + letter,
                                      word_node.children[letter], word_list)

And putting it altogether:

class Node:

    def __init__(self, character):
        self.character = character if character is None else character.lower()
        self.children = dict()
        self.terminus = False

    def add(self, child_node):
        self.children[child_node.character] = child_node


class Trie:

    def __init__(self):
        self._root = Node(None) 

    def insert(self, word):
        if not word: 
            return
        self.__insert(self._normalize_word(word), self._root).terminus = True

    def __insert(self, word, node):
        if not word:
            return node
        elif {word[0]} & node.children.keys():
            return self.__insert(word[1:], node.children[word[0]])
        else:
            new_node = Node(word[0])
            node.add(new_node)
            return self.__insert(word[1:], new_node)

    def __contains(self, word, node):
        if not word:
            return node
        return self.__contains(word[1:], node.children[word[0]]) if \
               {word[0]} & node.children.keys() else None

    def __contains__(self, item):
        word = self._normalize_word(item)
        return self.__contains(word, self._root) is not None

    def _normalize_word(self, word):
        return word.strip().lower()

    def get_possible_words(self, word):
        word_node = self.__contains(word, self._root)
        return [] if word_node is None else \
               self.__get_possible_words(word, word_node)

    def __get_possible_words(self, word, word_node, word_list=None):
        word_list = [] if word_list is None else word_list
        if word_node.terminus:
            word_list.append(word)
        for letter in word_node.children:
            if not word_node.children[letter]:
                word_list.append(word + letter)
            else:
                self.__get_possible_words(word + letter,
                                          word_node.children[letter], word_list)
        return word_list

    def get_all_words(self):
        word_list = []
        for letter in self._root.children:
            word_list.extend(self.get_possible_words(letter))
        return word_list