211. 添加与搜索单词 数据结构设计(trie 树)

 

labuladong 题解

难度中等429

请你设计一个数据结构，支持 添加新单词 和 查找字符串是否与任何先前添加的字符串匹配 。

实现词典类 WordDictionary ：

• WordDictionary() 初始化词典对象
• void addWord(word) 将 word 添加到数据结构中，之后可以对它进行匹配
• bool search(word) 如果数据结构中存在字符串与 word 匹配，则返回 true ；否则，返回  false 。word 中可能包含一些 '.' ，每个 . 都可以表示任何一个字母。

示例：

输入：
["WordDictionary","addWord","addWord","addWord","search","search","search","search"]
[[],["bad"],["dad"],["mad"],["pad"],["bad"],[".ad"],["b.."]]
输出：
[null,null,null,null,false,true,true,true]

解释：
WordDictionary wordDictionary = new WordDictionary();
wordDictionary.addWord("bad");
wordDictionary.addWord("dad");
wordDictionary.addWord("mad");
wordDictionary.search("pad"); // 返回 False
wordDictionary.search("bad"); // 返回 True
wordDictionary.search(".ad"); // 返回 True
wordDictionary.search("b.."); // 返回 True

struct TrieNode {
    bool is_end;
    std::vector<TrieNode*> children;
    TrieNode() {
        is_end = false;
        children = std::vector<TrieNode*>(26);
    }
};

class WordDictionary {
    TrieNode* root;
public:
    WordDictionary() {
        root = new TrieNode();
    }
    
    void addWord(string word) {
        TrieNode* node = root;
        for(auto ch : word) {
            if (node->children[ch-'a'] == nullptr) {
                node->children[ch-'a'] = new TrieNode();
            }
            node = node->children[ch - 'a'];
        }
        node->is_end = true;
    }

    bool dfs(string word, int index, TrieNode* root) {
        if (index >= word.size()){
            return root->is_end;
        }
        auto ch = word[index];
        
        if (ch == '.') {
            for(int i = 0; i < 26;i++) {
                TrieNode * node = root->children[i];
                if (node!=nullptr && dfs(word,index+1,node)) {
                    return true;
                }
            }
        } else if ('a' <= ch && ch <='z') {
            TrieNode* node = root->children[ch-'a'];
            if (node!=nullptr && dfs(word,index+1,node)) {
                return true;
            }
        }

        return false;
    }
        
    bool search(string word) {
        return dfs(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);
 */