Trie树python实现

class TrieNode:
    def __init__(self):
        self.children = {}
        self.flag = None
class Trie:
    def __init__(self):
        self.root = TrieNode()
    def add(self,s):
        curNode = self.root
        for i in s:
            if curNode.children.get(i) is None:
                newNode = TrieNode()
                curNode.children[i] = newNode
            curNode = curNode.children[i]
        curNode.flag = 1
    def search_node(self,s):
        curNode = self.root
        for i in s:
            if curNode.children.get(i) is None:
                return None
            else:
                curNode = curNode.children.get(i)
        if curNode.flag == 1:
            return curNode
    def search(self,s):
        x = self.search_node(s)
        return s if x is not None else None
    def remove(self,s):
        x = self.search_node(s)
        if x is not None:
            x.flag = None

if __name__ == '__main__':
    x = Trie()
    x.add('str')
    x.add('str1')
    x.remove('str')
    print x.search('str1')