Leetcode May Challenge - 05/14: Implement Trie (Prefix Tree)(Python)

题目描述

Implement a trie with insert, search, and startsWith methods.
Note:
You may assume that all inputs are consist of lowercase letters a-z.
All inputs are guaranteed to be non-empty strings.

例子

Trie trie = new Trie();

trie.insert(“apple”);
trie.search(“apple”); // returns true
trie.search(“app”); // returns false
trie.startsWith(“app”); // returns true
trie.insert(“app”);
trie.search(“app”); // returns true

解释

实现一个前缀树。支持插入，搜索单词，搜索前缀的函数。输入端所有字母均为小写字母。

思路

Trie是什么这里不解释了，思路就是用字典存当前节点的所有孩子节点，并用一个布尔变量判断该节点是否为一个单词的结束。

代码

import collections

class Node(object):
    def __init__(self):
        self.child = collections.defaultdict(Node)
        self.end = False

class Trie(object):

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

    def insert(self, word):
        """
        Inserts a word into the trie.
        :type word: str
        :rtype: None
        """
        cur = self.root
        for w in word:
            cur = cur.child[w]
        cur.end = True
        

    def search(self, word):
        """
        Returns if the word is in the trie.
        :type word: str
        :rtype: bool
        """
        cur = self.root
        for w in word:
            if not cur.child.get(w):
                return False
            cur = cur.child[w]
        return cur.end
        

    def startsWith(self, prefix):
        """
        Returns if there is any word in the trie that starts with the given prefix.
        :type prefix: str
        :rtype: bool
        """
        cur = self.root
        for p in prefix:
            if not cur.child.get(p):
                return False
            cur = cur.child[p]
        return True