Add and Search Word - Data structure Design

题目描述

解题方法

Solution

from collections import defaultdict

class TrieNode(object):
    def __init__(self):
        self.is_word = False
        self.children = defaultdict(TrieNode)

class WordDictionary(object):
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.root = TrieNode()

    def addWord(self, word):
        """
        Adds a word into the data structure.
        :type word: str
        :rtype: void
        """
        current = self.root
        for letter in word:
            current = current.children[letter]
        current.is_word = True

    def search(self, word):
        """
        Returns if the word is in the data structure. A word could
        contain the dot character '.' to represent any one letter.
        :type word: str
        :rtype: bool
        """
        current = self.root
        return self.search_from(current, word)

    def search_from(self, current, word):
        if len(word) == 0:
            return current.is_word
        else:
            if word[0] == ".":
                for letter in current.children:
                    if self.search_from(current.children[letter], word[1:]):
                        return True
                return False
            else:
                if word[0] in current.children:
                    return self.search_from(current.children[word[0]], word[1:])
                else:
                    return False

Reference