Datová struktura trie (prefixový strom) v Pythonu

Datová struktura trie (prefixový strom) v Pythonu

Trie neboli prefixový strom je datová struktura sloužící k uložení asociativního pole. K jeho prvkům se přistupuje pomocí klíče, který je ve většině případů tvořen textovým řetězcem. Díky své vnitřní struktuře umožňuje trie rychlé vyhledání s lineární časovou složitostí. Trie bývá často využívána pro uložení slovníků, kde se kromě rychlosti vyhledání využívá ještě relativně příznivých nároků na paměť v případě, že se ve slovníku nachází velké množství slov se stejným prefixem.
Oproti hashovací tabulce  má trie výhodu v tom, že odpadá potřeba hledání optimální hashovací funkce a současně v případě rozdílných klíčů není třeba řešit kolize. Trie je ideální pro případy, kdy je třeba nad určitou množinou slov hledat všechna slova začínající určitým prefixem.

Příklad jak může trie vypadat

Představte si, že máte slova: “ok, oko, okolo, okap, osa” a chcete je uložit do datové struktury trie. Výsledek by mohl vypadat nějak takto:

Kořenový uzel
/
o  
/   
k*  s  
/      
o*  a   a* 
/   /        
l   p*         
/                
o*

Jak vidíte, trie je strom jehož uzly jsou písmena slov v něm uložených. Každý uzel si s sebou nese informaci, zda je koncovým písmenem nějakého slova. Protože uložená slova mají víceméně shodný prefix, stačilo nám k uložení pěti slov s celkovou délkou sedmnácti znaků jenom devět uzlů plus kořenový uzel.

Pár slov k mojí implementaci

Obecně

Program přijímá jméno textového souboru se slovníkem, z kterého bude trie sestavena, jako svůj první a jediný argument.
Při implementaci jsem využil kontejnerový typ defaultdict ze standardní knihovny Collections. Je to typ podobný klasickému slovníku dictionary. Jeho výhodou oproti dictionary je však možnost nastavení hodnoty, kterou má slovník vracet v případě, že se v něm klíč nenachází bez potřeby explicitního volání funkce setdefault.

Uzel trie (třída _Node)

Možností jak implementovat trii je několik. Já jsem se ji rozhodl implementovat tak, že každý uzel stromu má dvě proměnné:

  • Slovník obsahující ukazatele na synovské uzly, jehož klíče jsou znaky reprezentující tyto uzly.
  • Boolovskou hodnotu True, nebo False v závislosti na tom, zda posloupnost znaků která vedla k tomuto uzlu je slovo, nebo ne.

Třída má metodu pro přidání synovského uzlu add_char, metodu pro označení uzlu jako konec slova set_end, metodu pro test, zda je uzel konec slova is_end a metodu __getitem__, která umožňuje přístup k potomkům uzlu pomocí indexovacího operátoru hranaté závorky [ ].

Samotná trie (třída Trie)

Obsahuje tyto instanční proměnné:

  • Kořenový uzel stromu
  • Hledané slovo
  • Délku hledaného slova

Kromě konstruktoru, který přijímá jako parametr soubor se slovníkem, tu najdeme metodu __getitem__ sloužící pro vyhledání slova v trii pomocí indexovacího operátoru [ ].
Neveřejná metoda _check_word realizuje samotné rekurzivní vyhledání slova v trii a je volána metodou __getitem__.

#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
Trie datastructure (prefix tree).
Author: Tomas "Tojaj" Mlcoch
https://www.tojaj.com
"""

import sys
from collections import defaultdict

PROMPT = ':> '

class _Node(object):
    """
    Node of the prefix tree
    """

    __slots__ = ('__childrens', '__end')

    def __init__(self):

        self.__childrens = defaultdict(lambda: 0)
        self.__end = False

    def add_char(self, character):
        """ Add character to node """

        item = self.__childrens[character]
        if item == 0:
            item = _Node()
            self.__childrens[character] = item

        return item

    def __getitem__(self, character):

        if len(character) != 1:
            raise AttributeError('You have to use only one character!')

        if character in self.__childrens:
            return self.__childrens[character]
        else:
            return 0

    def set_end(self):
        """ Mark node as end of word """
        self.__end = True

    def is_end(self):
        """ Test if node is end of word """
        return self.__end

class Trie(object):
    """
    Trie data structure
    """

    __slots__ = ('__root', '__pattern', '__plen')

    def __init__(self, in_file):

        self.__root = _Node()
        self.__pattern = ''
        self.__plen = 0

        for line in file(in_file):
            tmp_root = self.__root
            for char in line.strip():
                tmp_root = tmp_root.add_char(char)
            tmp_root.set_end()

    def __getitem__(self, sword):

        self.__plen = len(sword)
        self.__pattern = sword + '_'

        if self.__plen == 0:
            return

        return self._check_word(self.__root)

    def _check_word(self, root, i=0):
        """ Check if word is in trie """

        if isinstance(root, int) or i > self.__plen:
            return False

        node = root[self.__pattern[i]]

        # End of pattern
        if i >= (self.__plen - 1):
            if (node != 0 and node.is_end()):
                return True
            else:
                return False

        if self._check_word(node, i+1):
            return True

        return False

if __name__ == '__main__':
    trie = Trie(sys.argv[1])  
    word = raw_input(PROMPT)
    while word != '':
        print trie[word.lower()]
        word = raw_input(PROMPT)
Comments are closed.