字典树简单实现

#coding=utf8

class TrieNode(object):
    def __init__(self):
        self.data = None
        self.is_word = False
        self.children = []
    
class Trie(object):
    def __init__(self):
        self.root = TrieNode()

    def add(self, s):
        p = self.root
        n = len(s)
        for i in range(n):
            if s[i] not in [ ele.data for ele in p.children]:
                new_node = TrieNode()
                new_node.data = s[i]
                if i == n - 1:
                    new_node.is_word = True
                p.children.append(new_node)
                p = new_node
            else:
                for ele in p.children:
                    if ele.data == s[i]:
                        if i == n-1:
                            ele.is_word = True
                        else:
                            p = ele
                            break

    def search(self,s):
        p = self.root
        n = len(s)
        for i in range(len(s)):
            if s[i] not in [ ele.data for ele in p.children]:
                return False
            else:
                for ele in p.children:
                    if i == n - 1 and ele.data == s[i]:
                        if ele.is_word == True:
                            return True
                        else:
                            return False

                    if ele.data == s[i]:
                        p = ele
                        break


                        
        
if __name__ == '__main__':
    trie = Trie()
    trie.add('str')
    trie.add('acb')
    trie.add('acblde')
    print trie.search('str')
    print trie.search('acb')
    print trie.search('ac')
    trie.add('ac')
    print trie.search('ac')



    """"
test:
    trie = Trie()
    trie.add('str')
    trie.add('acb')
    trie.add('acblde')
    print trie.search('str')
    print trie.search('acb')
    print trie.search('ac')
    trie.add('ac')
    print trie.search('ac')
print:
    >>>
    True
    True
    False
    True
    >>>

    """