# coding:utf-8 """ Implement a trie with insert, search, and startsWith methods. Note: You may assume that all inputs are consist of lowercase letters a-z. Subscribe to see which companies asked this question """
classTrieNode(object): def__init__(self): """ Initialize your data structure here. """ self.data = {} self.is_word = False
definsert(self, word): """ Inserts a word into the trie. :type word: str :rtype: void """ node = self.root for letter in word: child = node.data.get(letter) ifnot child: node.data[letter] = TrieNode() node = node.data[letter] node.is_word = True
defsearch(self, word): """ Returns if the word is in the trie. :type word: str :rtype: bool """ node = self.root for letter in word: node = node.data.get(letter) ifnot node: returnFalse return node.is_word # 判断单词是否是完整的存在在trie树中
defstarts_with(self, prefix): """ Returns if there is any word in the trie that starts with the given prefix. :type prefix: str :rtype: bool """ node = self.root for letter in prefix: node = node.data.get(letter) ifnot node: returnFalse returnTrue
defget_start(self, prefix): """ Returns words started with prefix :param prefix: :return: words (list) """ def_get_key(pre, pre_node): words_list = [] if pre_node.is_word: words_list.append(pre) for x in pre_node.data.keys(): words_list.extend(_get_key(pre + str(x), pre_node.data.get(x))) return words_list
words = [] ifnot self.starts_with(prefix): return words if self.search(prefix): words.append(prefix) return words node = self.root for letter in prefix: node = node.data.get(letter) return _get_key(prefix, node)
# Your Trie object will be instantiated and called as such: trie = Trie() trie.insert("somestring") trie.insert("somebody") trie.insert("somebody1") trie.insert("somebody3") print trie.search("key") print trie.search("somebody3") print trie.get_start('some')