Trie with insert, search, and startsWith methods for lowercase strings

I wrote the following prefix tree data structure for this Leet code challenge.


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


Note:

You may assume that all inputs are consist of lowercase letters a-z.

However it keeps timing out on the long input. As far as I know the insert, search and startsWith functions have \$\mathcal{O}(n)\$ time complexity, where \$n\$ is the length of the input string. Can someone review my code and see where I could improve the speed of the code?

```
class TrieNode(object):
def __init__(self):
self.next = [['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],
['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],
['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None],['!',None]]
self.isWord = False

class Trie(object):
def __init__(self):
self.root = TrieNode()

def insert(self, word):
"""
Inserts a word into the trie.
:type word: str
:rtype: void
"""
i = 0
wordLength = len(word)
node = self.root
while i < wordLength:
idx = ord(word[i]) - ord('a')
if node.next[idx][0] == '!':
node.next[idx][0] = word[i]
if i < wordLength-1 and node.next[idx][1] == None:
node.next[idx][1] = TrieNode()
elif i == wordLength-1:
node.isWord = True
return
i += 1
node = node.next[idx][1]

def search(self, word):
"""
Returns if the word is in the trie.
:type word: str
:rtype: bool
"""
i = 0
wordLength = len(word)
node = self.root
while i < wordLength and node != None:
idx = ord(word[i]) - ord('a')
if node.next[idx][0] == '!':

Assuming the next attribute of TrieNode is just a lookup for letters of the alphabet, you might be better off replacing it with a dictionary that maps letters of the alphabet to their next value. Something like:

class TrieNode(object):
    def __init__(self):
        self.next = {letter: ['!', None] for letter in string.ascii_lowercase}
        self.isWord = False

and now, rather than mucking around with ord() in your insert() and search() methods, you can use the characters directly. That won’t fundamentally change the complexity, but will make a small performance saving.

I think insert and search of a trie is O(key_length), so this is probably the sort of saving you need. You can’t change the complexity.

A few other small suggestions:

• Comparisons to None should use foo is None and foo is not None, not direct equality.

-
Once you’re using the characters as keys in the next attribute, you’d be better of replacing those while i

-
In general, PEP 8 convention is that Python variables are lowercase with underscores (e.g. is_word, starts_with), and only classes are camel case. (Although as @Graipher notes in the comments, camel case is specified for startsWith()`.)

StackExchange Code Review Q#143016, answer score: 2