Problem Description:
Implement a trie with insert
, search
, and startsWith
methods.
Note:
You may assume that all inputs are consist of lowercase letters a-z
.
这是一道小的设计题,设计一个字典树节点的数据结构,从而实现字典树的插入,查找字符串和查找前缀。
由于题目假设所有的输入都是a-z的26个字母,这道题目就可以直接用数组来表示每个节点的26个孩子指针(next)。另外,为了区分完整的字符串和字符串前缀,我们需要在每个节点处标记该处是不是字符串的结尾,本来一个bool就够用了,我这里用的是int表示的cnt,这样还可以实现字符串的计数功能。
有了字典树的结构之后,那三个操作都是很显然的事情,这里不再赘述。
下面简单分析该方法的优缺点:
优点,虽然用数组来存储,但是可以直接用字符来寻址(孩子节点的位置),相当于hash。上述三种操作的时间复杂度都是O(k),k是输入字符串的长度。
缺点,内存消耗特别大,每层节点数目按照26倍进行指数增长,处理的字符串肯定不能太长。这里还有一个小的改进可以稍微减少一点空间的消耗。细看我的实现,可以发现当插入一个字符时,不管它是否还有孩子节点,都直接预先分配了26个空的孩子节点指针。我们可以等某个节点有了子孩子之后,再一次性分配内存。
class TrieNode {
public:
// Initialize your data structure here.
TrieNode() {
cnt = 0;
next.assign(26, NULL);
}
int cnt;
vector<TrieNode*> next;
};
class Trie {
public:
Trie() {
root = new TrieNode();
}
// Inserts a word into the trie.
void insert(string word) {
TrieNode* ptr = root;
for(int i = 0; i < word.length(); ++i)
{
int idx = word[i] - 'a';
if(ptr->next[idx] == NULL)
{
ptr->next[idx] = new TrieNode();
}
ptr = ptr->next[idx];
}
++(ptr->cnt);
}
// Returns if the word is in the trie.
bool search(string word) {
TrieNode* ptr = root;
for(int i = 0; i < word.length(); ++i)
{
int idx = word[i] - 'a';
if(ptr->next[idx] == NULL)
{
return false;
}
ptr = ptr->next[idx];
}
if(ptr->cnt > 0) return true;
else return false;
}
// Returns if there is any word in the trie
// that starts with the given prefix.
bool startsWith(string prefix) {
TrieNode* ptr = root;
for(int i = 0; i < prefix.length(); ++i)
{
int idx = prefix[i] - 'a';
if(ptr->next[idx] == NULL)
{
return false;
}
ptr = ptr->next[idx];
}
return true;
}
private:
TrieNode* root;
};