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 &lt; wordLength:
            idx = ord(word[i]) - ord('a')
            if node.next[idx][0] == '!':
                node.next[idx][0] = word[i]
            if i &lt; 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 &lt; wordLength and node != None:
            idx = ord(word[i]) - ord('a')
            if node.next[idx][0] == '!':
                return False
            if i == wordLength - 1 and node.next[idx][1] != None:
                return node.isWord
            node = node.next[idx][1]
            i += 1
        if node == None and i &lt; wordLength:
            return False
        return True

    def startsWith(self, prefix):
        """
        Returns if there is any word in the trie
        that starts with the given prefix.
        :type prefix: str
        :rtype: bool
        """
        i = 0
        wordLength = len(prefix)
        node = self.root
        while i &lt; wordLength and node != None:
            idx = ord(prefix[i]) - ord('a')
            if node.next[idx][0] == '!':
                return False
            node = node.next[idx][1]
            i+=1
        if node == None and i &lt; wordLength:
            return False
        return True
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 &lt; wordLength loops with something like for i, char in enumerate(word).
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().)