A post about Haskell vs. Python readability came onto my radar the other day. It compares the implementation of a trie structure, and after looking upon the Python version I wanted to make my own attempt. I didn't make it to necessarily compare or "battle" against the other solutions, it's more of an exercise in the vein of "how would I do it".

The code

Here's the original code (for easier lookup, as I refer to a few things in it in the notes):

class Trie(object):
    def __init__(self, value=None):
        self.children = {}
        self.value = value
        self.flag = False # Flag to represent that a word ends at this node

    def add(self, char):
        val = self.value + char if self.value else char
        self.children[char] = Trie(val)

    def insert(self, word):
        node = self
        for char in word:
            if char not in node.children:
                node.add(char)
            node = node.children[char]
        node.flag = True

    def find(self, word):
        node = self
        for char in word:
            if char not in node.children:
                return None
            node = node.children[char]
        return node.value

    def all_prefixes(self, wlist):
        results = set()
        if self.flag:
            results.add(self.value)
        if not self.children: return results
        return reduce(lambda a, b: a | b,
                     [node.all_prefixes() for
                      node in self.children.values()]) | results

    def autocomplete(self, prefix):
        node = self
        for char in prefix:
            if char not in node.children:
                return set()
            node = node.children[char]
        return node.all_prefixes()

My code:

class Trie:
    def __init__(self):
        self.children = {}
        self.is_word_end = False

    def insert(self, word):
        for char in word:
            self = self.children.setdefault(char, Trie())
        self.is_word_end = True

    def words_with(self, prefix):
        if self.is_word_end:
            yield prefix
        for char, node in self.children.items():
            yield from node.words_with(prefix + char)

    def autocomplete(self, prefix):
        try:
            for char in prefix:
                self = self.children[char]
            return list(self.words_with(prefix))
        except KeyError:
            return []

A few notes:

P.S. Please don't talk to me about types and dataclasses, I will ridicule you to no end :-)

Comments: 3

  1. Andre Müller

    P.S. Please don't talk to me about types and dataclasses, I will ridicule you to no end :-)

    from dataclasses import dataclass, field
    
    
    @dataclass
    class Trie:
        children: dict = field(default_factory=dict, repr=False)
        is_word_end: bool = field(default=False)
    
        def insert(self, word):
            for char in word:
                self = self.children.setdefault(char, Trie())
            self.is_word_end = True
    
        def words_with(self, prefix):
            if self.is_word_end:
                yield prefix
            for char, node in self.children.items():
                yield from node.words_with(prefix + char)
    
        def autocomplete(self, prefix):
            try:
                for char in prefix:
                    self = self.children[char]
                return list(self.words_with(prefix))
            except KeyError:
                return []
    
    
    t = Trie()
    print(t)
    

    What are the disadvantages of this?

  2. Ivan Sagalaev

    What are the disadvantages of this?

    There are no advantages either. And field(default_factory=dict, repr=False) looks pretty messy, to boot.

  3. Nerio

    Thank you... Oh God, thank you, Python has no types Python needs no types (of course it has types, but you know, LOTR)

Add comment