LeetCode Add and Search Word - Data structure design

Add and Search Word - Data structure design 

Design a data structure that supports the following two operations:

void addWord(word)
bool search(word)

search(word) can search a literal word or a regular expression string containing only letters a-z or .. A . means it can represent any one letter.

For example:

addWord("bad")
addWord("dad")
addWord("mad")
search("pad") -> false
search("bad") -> true
search(".ad") -> true
search("b..") -> true

Note:
You may assume that all words are consist of lowercase letters a-z.

这个题目是字典树的扩展，之前有个 构建字典树的题目，此题知识按照需求增加一次搜索。

class TrieNode {
public:
	// Initialize your data structure here.
	TrieNode() 
	{
		memset(branch, 0, sizeof(branch));
		isWord = false;
	}

	TrieNode* branch[26];//记录26个小写字母的分支
	bool isWord;//判断某个分支节点是不是单词（比如插入abcd，只有abcd是单词，而abc只是前缀，不是单词；用于辅助search方法）
};

class WordDictionary {
public:
	//构造函数
	WordDictionary()
	{
		root = new TrieNode();
	}

	// Adds a word into the data structure.
	void addWord(string word)
	{
		TrieNode* p = root;
		for (int i = 0; i < word.length(); i++)
		{
			if (p->branch[word[i] - 'a'] == NULL)
			{
				p->branch[word[i] - 'a'] = new TrieNode();
			}

			p = p->branch[word[i] - 'a'];
		}
		//插入“是单词”的标记
		p->isWord = 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 MySearch(word, root);
	}

	bool MySearch(string word, TrieNode* p)
	{
		if (word.empty())
			return p->isWord;

		int index = word[0] - 'a';

		if (word[0] == '.')
		{
			for (int i = 0; i < 26; i++)
			{
				if (p->branch[i] == NULL)
					continue;
				//递归，搜索子串，分支
				if (MySearch(word.substr(1, word.size() - 1), p->branch[i]))
					return true;
			}
			return false;
		}
		else if (p->branch[index] != NULL)
		{
			if (word.size() > 1)
				return MySearch(word.substr(1, word.size() - 1), p->branch[index]);
			else
				return p->branch[index]->isWord;
		}

		return false;
	}

private:
	TrieNode* root;
};