Python Trie树实现最长前缀字符串提取

在文本解析项目中,经常会碰到提取品牌、商家名等需求。如给定一个手机型号字符串,要求从中提取出品牌。Trie可以很好满足此类需求。

Tire,也叫前缀树字典树,是一种数据结构,可以用来快速检索字符串是否存在以及在字符串开始处抽取预定义的子字符串。搜索时间复杂度为O(M)  M为字符串长度。

Python Trie树实现最长前缀字符串提取_第1张图片 Tire

代码实现

Python中无指针,使用Dict实现树结构。

# -*- coding: utf-8 -*-
"""
Trie for prefix search, a data structure that quickly matches and extracts predefined substrings
at the beginning of a given text (if they can be found).

We can also skip certain characters and still succeed in a match.
"""

default_ignored_chars = u' _-/'


class Trie(object):
    def __init__(self, items, ignored_chars=default_ignored_chars):
        """ Stores all given items into this trie. """
        self.ignored_chars = ignored_chars

        self.trie = {}
        for item in items:
            assert item, 'Empty/none item passed in'
            item = item.strip()
            assert item, 'Empty item given'
            curr_dict = self.trie
            for c in item.upper():
                if c not in self.ignored_chars:
                    curr_dict = curr_dict.setdefault(c, {})
            curr_dict['end'] = item

    def is_item(self, text):
        """ Return True if text is a valid item stored in this trie. """
        if not text:
            return False
        curr_dict = self.trie
        for c in text.upper():
            if c not in self.ignored_chars:
                if c not in curr_dict:
                    return False
                curr_dict = curr_dict[c]
        return 'end' in curr_dict

    def extract_longest_item(self, text):
        """ Return longest item-name found at beginning of the text. Also returns the
            offset where the item ends in case the caller wants to chop the string. """
        curr_dict, longest, offset = self.trie, None, 0

        if not text:
            return longest, offset

        for i, c in enumerate(text.upper()):
            if c not in self.ignored_chars:
                if c not in curr_dict:
                    return longest, offset
                curr_dict = curr_dict[c]
                if 'end' in curr_dict:
                    longest, offset = curr_dict['end'], i + 1

        return longest, offset


# tester
if __name__ == '__main__':

    brands = ['Huawei', 'OPPO', 'VIVO', 'Xiaomi', 'Xiao', 'HTC', 'Oneplus']
    model_name = 'xiaomi mix3'
    brand_lookup = Trie(brands)

    brand, offset = brand_lookup.extract_longest_item(model_name)
    print(brand, offset)

通过自己写代码理解Trie之后,也可以使用Google的开源实现  https://github.com/google/pygtrie

你可能感兴趣的:(Python)