Trie

Problem

208. Implement Trie (Prefix Tree)

A simple code (140 ms)

class Trie:
def __init__(self):
        """
        Initialize your data structure here.
        """
        self._trie = {}
        self._END = '_END_'
def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        cur_d = self._trie
        for w in word:
            cur_d = cur_d.setdefault(w, {})
        cur_d[self._END] = self._END
def search(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        cur_d = self._trie
        for w in word:
            if w in cur_d:
                cur_d = cur_d[w]
            else:return False
        return self._END in cur_d
        
    def startsWith(self, prefix: str) -> bool:
        """
        Returns if there is any word in the trie that starts with the given prefix.
        """
        cur_d = self._trie
        for w in prefix:
            if w in cur_d:
                cur_d = cur_d[w]
            else:return False
        return True

The critical part is the code of

cur_d = cur_d.setdefault(w, {})

It means if the current dict has a dictionary as value, we get that value. That value will be the next level dict. If the current dict does not contain the related key, we will generate a key:{} for that key.

>>> d = {}
>>> d
{}
>>> d['a'] = 1
>>> d
{'a': 1}
>>> d.setdefault('b', {})
{}
>>> d
{'a': 1, 'b': {}}
>>> d.setdefault('a',{})
1

A longer and slower design

The code is from here

class TrieNode:
    # Initialize your data structure here.
    def __init__(self):
        self.children = collections.defaultdict(TrieNode)
        self.is_word = False
class Trie:
def __init__(self):
        self.root = TrieNode()
def insert(self, word):
        current = self.root
        for letter in word:
            current = current.children[letter]
        current.is_word = True
def search(self, word):
        current = self.root
        for letter in word:
            current = current.children.get(letter)
            if current is None:
                return False
        return current.is_word
def startsWith(self, prefix):
        current = self.root
        for letter in prefix:
            current = current.children.get(letter)
            if current is None:
                return False
        return True

A standard C++ code

struct vertex 
{
    char alphabet;
    bool exist;
    vector<vertex*> child;
    vertex(char a): alphabet(a), exist(false) {child.assign(26, nullptr);}
};
class Trie {
private:
    vertex* root;
public:
    /** Initialize your data structure here. */
    Trie() 
    {
        root = new vertex('!');
    }
    
    /** Inserts a word into the trie. */
    void insert(string word) 
    {
        vertex* curr = root;
        for (int i = 0; i < word.size(); ++i)
        {
            int alpha_num = word[i] - 'a';
            // add a node if null
            if (curr->child[alpha_num] == nullptr)
            {
                curr->child[alpha_num] = new vertex(word[i]);
            }
            curr = curr->child[alpha_num];
        }
        curr->exist = true;
    }
    
    /** Returns if the word is in the trie. */
    bool search(string word) 
    {
        vertex* curr = root;
        for (int i = 0; i < word.size(); ++i)
        {
            int alpha_num = word[i] - 'a';
            // add a node if null
            if (curr->child[alpha_num] == nullptr)
            {
               return false;
            }
            curr = curr->child[alpha_num];
        }
        return curr->exist;  
    }
    
    /** Returns if there is any word in the trie that starts with the given prefix. */
    bool startsWith(string prefix) 
    {
        vertex* curr = root;
        for (int i = 0; i < prefix.size(); ++i)
        {
            int alpha_num = prefix[i] - 'a';
            // add a node if null
            if (curr->child[alpha_num] == nullptr)
            {
               return false;
            }
            curr = curr->child[alpha_num];
        }
        return true;
        
    }
};
/**
 * Your Trie object will be instantiated and called as such:
 * Trie* obj = new Trie();
 * obj->insert(word);
 * bool param_2 = obj->search(word);
 * bool param_3 = obj->startsWith(prefix);
 */

211. Add and Search Word — Data structure design

class vertex 
{
    public:
    char alphabet;
    bool exist;
    vector<vertex*> child;
    vertex(char a): alphabet(a), exist(false) {child.assign(26, nullptr);}
};
class WordDictionary {
private:
    vertex* root;
     bool search(string word, int i, vertex* curr)
    {
        if (word[i] != '.')
        {
            int alpha_num = word[i] - 'a';
            if (curr->child[alpha_num] == nullptr) return false;
            else 
            {
                if (i == word.length() - 1) return curr->child[alpha_num]->exist;
                curr = curr->child[alpha_num];
                return search(word, i+1, curr);
            }
        }
        else
        {
            for (int j = 0; j < 26; ++j)
            {
                if (curr->child[j] == nullptr) continue;
                else
                {
                    if (i == word.length() - 1)
                    {
                        if (curr->child[j]->exist) return true;
                        else continue;
                    }
                    else if (search(word, i+1, curr->child[j])) return true;
                }
            }
            return false;
        }
    }
public:
    /** Initialize your data structure here. */
    WordDictionary() {
        root = new vertex('!');
    }
    
    /** Adds a word into the data structure. */
    void addWord(string word) {
        vertex* curr = root;
        for (int i = 0; i < word.size(); ++i)
        {
            int alpha_num = word[i] - 'a';
            // add a node if null
            if (curr->child[alpha_num] == nullptr)
            {
                curr->child[alpha_num] = new vertex(word[i]);
            }
            curr = curr->child[alpha_num];
        }
        curr->exist = true;
    }
    
    /** Returns if the word is in the data structure. A word could contain the dot character '.' to represent any one letter. */
    bool search(string word) 
    {
       return search(word, 0, root);
        
    }
};
/**
 * Your WordDictionary object will be instantiated and called as such:
 * WordDictionary* obj = new WordDictionary();
 * obj->addWord(word);
 * bool param_2 = obj->search(word);
 */