8
\$\begingroup\$

Question:

Implement a trie with insert, search, and startsWith methods.

Source:

https://leetcode.com/problems/implement-trie-prefix-tree/description/

The Trie object will be instantiated and called as such:

obj = Trie()
obj.insert(word)
param_2 = obj.search(word)
param_3 = obj.startsWith(prefix)

Code: I have a TrieNode class and a Trie class. I'm not sure if I can make my code more pythonic. I ran the code on leetcode and it passed all the testcases.

class TrieNode:

    def __init__(self):
        self.children = [0]*26
        self.end = False

    def put(self, ch):
        self.children[ord(ch)-ord('a')] = TrieNode()

    def get(self, ch):
        return self.children[ord(ch)-ord('a')]

    def contains_key(self, ch):
        return self.children[ord(ch)-ord('a')] != 0

    def set_end(self):
        self.end = True

    def is_end(self):
        return self.end


class Trie:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = TrieNode()


    def insert(self, word):
        """
        Inserts a word into the trie.
        :type word: str
        :rtype: void
        """
        curr = self.root
        for i in range(len(word)):
            if not curr.contains_key(word[i]):
                curr.put(word[i])
            curr = curr.get(word[i])

        curr.set_end()


    def search(self, word):
        """
        Returns if the word is in the trie.
        :type word: str
        :rtype: bool
        """
        curr = self.root
        for i in range(len(word)):
            if not curr.contains_key(word[i]):
                return False
            curr = curr.get(word[i])

        return True if curr.is_end() else False


    def startsWith(self, prefix):
        """
        Returns if there is any word in the trie that starts with the given prefix.
        :type prefix: str
        :rtype: bool
        """
        curr = self.root
        for i in range(len(prefix)):
            if not curr.contains_key(prefix[i]):
                return False
            curr = curr.get(prefix[i])

        return True
\$\endgroup\$
0

1 Answer 1

Reset to default
4
\$\begingroup\$

for with range and len

This is one of the most common things to be seen in python code to someone that isn't very familiar to python, however it's not Pythonic and almost always the wrong way to do it.

Firstly if you are not using the index, so there is no need to iterate base on the index. Consequently this can be replaced with a normal for over a collection, in your case, a string.

So this:

for i in range(len(word)):
    if not curr.contains_key(word[i]):
        curr.put(word[i])
    curr = curr.get(word[i])

Should become this:

for letter in word:
    if not curr.contains_key(letter):
        curr.put(letter)
    curr = curr.get(letter)

Note that the index was not necessary at all, and the code actually becomes clearer.

Even in cases where the index of the element is needed for something, you can use enumerate which gives you a tuple for each element, holding both the index and the element:

for index, letter in enumerate(word):

Which is still more readable and remains Pythonic

Keys stored in array

Even though it is reasonable to store keys in a array it is actually simpler to use a dictionary. Doing so avoids having to convert a character to value in the range 0-26 and still makes it simple to check whether it exists.

So instead of doing self.children = [0]*26 you would do self.children = {}. Most of the TrieNode methods would have to be rewritten accordingly:

def put(self, ch):
    self.children[ch] = TrieNode()

def get(self, ch):
    return self.children[ch]

def contains_key(self, ch):
    return ch in self.children

Note that they are all so small and simple that i would go as far as removing them and using curr.children directly in both in insert, search and startsWith.

i.e (insert):

curr = self.root
for letter in word:
    if letter not in curr.children:
        curr.children[letter] = TrieNode()
    curr = curr.children[letter]

While using curr.children makes sense in this simple example, it may not be the case if there is a lot more logic in TrieNode. Each case is unique in a way, and what may be fine in one scenario may be bad in another.

The same can be said about the is_end method that merely returns end. Following this last step, you could go a step further and also use end directly as there isn't much to encapsulate. On top of that only Trie uses TrieNode.

On a side note, in cases where you need a getter, you should use the @property decorator which would be far more appropriate and Pythonic.

camelCase

I know the methods stubs were provided in the problem description, however, camelCase is not the typical naming convention for python methods. The preferred style for this is snake_case, and as such, startsWith should be renamed to starts_with.

Feel free to take a look at the related PEP, PEP 8 -- Style Guide for Python Code

\$\endgroup\$

You must log in to answer this question.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.