211.md

211. Design Add and Search Words Data Structure

Solution 1

class TreeNode(object):
    def __init__(self):
        self.is_word = False
        self.children = {c: None for c in "abcdefghijklmnopqrstuvwxyz"}

class WordDictionary(object):

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = TreeNode()

    def addWord(self, word):
        """
        :type word: str
        :rtype: None
        """
        cur_node = self.root
        for c in word:
            if not cur_node.children[c]:
                cur_node.children[c] = TreeNode()
            cur_node = cur_node.children[c]
        cur_node.is_word = True

    def search(self, word):
        """
        :type word: str
        :rtype: bool
        """
        return self._search(word, self.root, 0)
    
    def _search(self, word, cur_node, i):
        if i == len(word):
            return cur_node.is_word
        if word[i] != ".":
            if not cur_node.children[word[i]]:
                return False
            else:
                return self._search(word, cur_node.children[word[i]], i + 1)
        else:
            for c in "abcdefghijklmnopqrstuvwxyz":
                if cur_node.children[c] and self._search(word, cur_node.children[c], i + 1):
                    return True
            return False


# Your WordDictionary object will be instantiated and called as such:
# obj = WordDictionary()
# obj.addWord(word)
# param_2 = obj.search(word)