常用python算法實現（二）——前綴樹TrieTree

一.概述（多個關鍵詞-實體檢索-查詢）

         TrieTree（前綴樹），又被稱爲字典樹、單詞查找樹，是一種比較常見的數據存儲結構與算法。

         顧名思義，前綴樹便是公共的字符只保存一次的多路樹。如你所見，它的基本思想是以時間換空間，時間複雜度爲logN，效果還不錯。不過，我覺得它應用廣泛的另外一個原因是它保存了字符的順序。

         應用: 字符串檢索、查詢與排序，前綴與公共前綴等。

        github地址: https://github.com/yongzhuo/Tookit-Sihui/blob/master/tookit_sihui/ml_common/trie_tree/trie_tree.py

二.實現（多個關鍵詞-實體檢索-查詢）

            實現的是: 從句中中查找並提取關鍵詞或者是實體

# -*- coding: UTF-8 -*-
# !/usr/bin/python
# @time     :2019/6/27 16:40
# @author   :Mo
# @function :TrieTree of keywords find, 只返回查全的情況, 查找句子中的關鍵詞（例如影視名、人名、關鍵詞、實體等）


import logging
logger = logging


class TrieNode:
    """
        前綴樹節點-鏈表
    """
    def __init__(self):
        self.child = {}


class TrieTree:
    """
        前綴樹構建, 新增關鍵詞, 關鍵詞詞語查找等
    """
    def __init__(self):
        self.algorithm = "trietree"
        self.root = TrieNode()

    def add_keyword(self, keyword):
        """
            新增一個關鍵詞
        :param keyword: str, 構建的關鍵詞
        :return: None
        """
        node_curr = self.root
        for word in keyword:
            if node_curr.child.get(word) is None:
                node_next = TrieNode()
                node_curr.child[word] = node_next
            node_curr = node_curr.child[word]
        # 每個關鍵詞詞後邊, 加入end標誌位
        if node_curr.child.get('[END]') is None:
            node_next = TrieNode()
            node_curr.child['[END]'] = node_next
        node_curr = node_curr.child['[END]']
        logger.info("add {} success!".format("".join(keyword)))

    def delete_keyword(self, keyword):
        """
            刪除一個關鍵詞
        :param keyword: str, 構建的關鍵詞
        :return: None
        """
        node_curr = self.root
        flag = 1
        for word in keyword:
            if node_curr.child.get(word) is not None:
                node_curr = node_curr.child[word]
            else:
                flag = 0
        # 每個關鍵詞詞後邊, 加入end標誌位
        if node_curr.child.get('[END]') is not None and flag == 1:
            node_curr.child.pop('[END]')
        else:
            logger.info("{} is not in trietree, delete keyword faild!".format("".join(keyword)))

    def add_keywords_from_list(self, keywords):
        """
            新增關鍵詞s, 格式爲list
        :param keyword: list, 構建的關鍵詞
        :return: None
        """
        for keyword in keywords:
            self.add_keyword(keyword)

    def find_keyword(self, sentence):
        """
            從句子中提取關鍵詞, 可提取多個
        :param sentence: str, 輸入的句子
        :return: list, 提取到的關鍵詞
        """
        assert type(sentence) == str
        if not sentence: # 空格字符不取
            return []

        node_curr = self.root # 關鍵詞的頭, 每遍歷完一遍後需要重新初始化
        index_last = len(sentence)
        keyword_list = []
        keyword = ''
        count = 0
        for word in sentence:
            count += 1
            if node_curr.child.get(word) is None: # 查看有無後綴, 即匹配到一個關鍵詞最後一個字符的時候
                if keyword: # 提取到的關鍵詞(也可能是前面的幾位)
                    if node_curr.child.get('[END]') is not None: # 取以end結尾的關鍵詞
                        keyword_list.append(keyword)
                    if self.root.child.get(word) is not None: # 處理連續的關鍵詞情況, 如"第九區流浪地球"
                        keyword = word
                        node_curr = self.root.child[word]
                    else: #
                        keyword = ''
                        node_curr = self.root  # 重新初始化
            else: # 有後綴就加到name裏邊
                keyword = keyword + word
                node_curr = node_curr.child[word]
                if count == index_last:  # 實體結尾的情況
                    if node_curr.child.get('[END]') is not None:
                        keyword_list.append(keyword)
        return keyword_list

    def match_keyword(self, keyword):
        """
            判斷keyword在不在trietree裏邊
        :param keyword: str, input word
        :return: boolean, True or False
        """
        node = self.root
        for kw in keyword:
            if not node.child.get(kw):
                return False
            node = node.child[kw]
        if not node.child.get('[END]'):
            return False
        return True


def get_trie_tree_class(keywords):
    """
        根據list關鍵詞，初始化trie樹
    :param keywords: list, input
    :return: objext, 返回實例化的trie
    """
    trie = TrieTree()
    trie.add_keywords_from_list(keywords)
    return trie


if __name__ == "__main__":
    print("".join("你好呀"))
    # 測試1, class實例
    trie = TrieTree()
    keywords = ['英雄', '人在囧途', '那些年,我們一起追過的女孩', '流浪地球', '華娛',
                '犬夜叉', '火影', '名偵探柯南', '約會大作戰', '名作之壁', '動漫',
                '乃木坂46', 'akb48', '飄', '最後的武士', '約會', '英雄2', '日娛',
                '2012', '第九區', '星球大戰', '侏羅紀公園', '泰坦尼克號', 'Speed']
    keywords = [list(keyword.strip()) for keyword in keywords]
    trie.add_keywords_from_list(keywords) # 創建樹
    keyword = trie.find_keyword('第九區約會, 侏羅紀公園和泰坦尼克號泰坦尼克號')
    print(keyword)
    gg = trie.delete_keyword('英雄')
    gg = trie.delete_keyword('英雄3')

    keyword = trie.match_keyword('英雄')
    keyword2 = trie.match_keyword('英雄2')

    print(keyword)


    # 測試2, get樹
    trie_tree = get_trie_tree_class(keywords) # 創建樹並返回實例化class
    while True:
        print("sihui請你輸入:")
        input_ques = input()
        keywords = trie_tree.find_keyword(input_ques)
        print(keywords)

希望對你有所幫助!