Problem

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

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);
*/

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);
*/

Data Scientist/MLE/SWE @takemobi