Доделал на днях питонью библиотеку datrie, реализующую префиксное дерево (см. википедию или хабр), спешу поделиться.
Если вкратце, то можно считать, что datrie.Trie — это замена стандартному питоньему dict, которая при определенных условиях (ключи — строки) занимает меньше памяти, имеет сравнимую скорость получения отдельного элемента и поддерживает дополнительные операции (получение всех префиксов данной строки, получение всех строк, начинающихся с данной строки и др.), которые работают примерно так же быстро, как и «словарные» операции.
Работает под Python 2.6-3.3, поддерживает юникод, лицензия LGPL.
datrie — это Cython-обертка над библиотекой libdatrie. В libdatrie реализован вариант префиксного дерева, в котором конечный автомат переходов между узлами хранится в специальных 2 массивах + реализовано сжатие суффиксов (ветки, которые не ветвятся, хранятся еще в одном отдельном массиве «хвостов»). Это не совсем «state of art» вариант trie (HAT-trie и тд вроде должны быть быстрее), но вариант достаточно быстрый/эффективный и с хорошей готовой реализацией (а реализация может на практике убить любой алгоритм).
Существующие варианты trie-образных структур для питона меня не устраивали. Чисто питоньи реализации неизбежно будут потреблять очень много памяти, они отметаются сразу. Другие реализации trie-образных структур для питона:
- trie.c в biopython. Не поддерживает юникод (хотя это ОК), не работает под 3.х (и чтоб заставить работать под 3.х, нужно хорошенько переписывать обертку, т.к. библиотека реализована как «голое» C-расширение);
- github.com/buriy/python-chartrie от buriy — довольно быстрая (на __getitem__ быстрее, чем datrie в итоге), занимает много памяти (чтоб это поправить, нужно, видимо, структуру данных менять, т.е. основу библиотеки), мало функционала, не поддерживает юникод (хотя это ОК);
- www.dalkescientific.com/Python/PyJudy.html — старая, сложная, непонятно, работает ли (написано, что поддерживает 2.3 и 2.4)
Может еще что-то и есть, но, в любом случае, варианта «поставил — заработало — всем устраивает!» не нашел, поэтому считаю, что желание вспомнить C и поразбираться получше в Cython и Python C-API вполне успешно удалось прикрыть отсутствием нормальной реализации trie для питона.
Установка (как это обычно бывает при установке расширений, потребуется компилятор):
pip install datrie
Создание:
import string
import datrie
trie = datrie.new(string.ascii_lowercase)
При создании нужно сразу сказать, какие ключи можно будет с этим trie использовать (явно указывать алфавит или диапазоны). Это ограничение libdatrie, которое позволяет эффективно хранить автомат состояний и поддерживать юникод — сначала указываются допустимые диапазоны unicode-символов, потом unicode-ключи транслируются в более компактное внутреннее представление.
Мне кажется, что можно было бы на практике обойтись однобайтовыми кодировками вроде cp1251 и иметь примерно тот же функционал и эффективность, но подход с диапазонами символов тоже работает неплохо, уж сделано так в libdatrie и ладно. Я поэтому пишу, что «не поддерживает юникод — это ОК» — в случае с trie однобайтовая кодировка может быть удобным вариантом.
Потом с trie можно работать как со словарем:
>>> trie[u'foo'] = 5
>>> trie[u'foobar'] = 10
>>> trie[u'bar'] = 'bar value'
>>> trie.setdefault(u'foobar', 15)
10
>>> u'foo' in trie
True
>>> trie[u'foo']
5
Ключи должны быть юникодными :) В питоне 3.х это вполне естественно, в 2.x в примерах приходится ставить букву u, уж простите. Значения могут быть любыми питоньими объектами, что нетипично для C-реализаций trie (обычно там целые число в качестве значений). На самом деле «внутри» значения и правда представляют собой целые числа, но datrie.Trie использует их как индексы в массиве «настоящих» значений. Для тех, кому эта фича не нужна (например, значения не интересуют совсем), в datrie есть datrie.BaseTrie, который немного быстрее и умеет хранить только числа.
Немного о скорости. Все замеры проводил на trie.Trie с сотней тысяч уникальных русский и английских слов (50/50) и int-значениями «1», другие замеры (на миллионе url'ов) можно глянуть тут или (еще лучше) проводить самим на своих данных. Замеры скорости я проводил только для того, чтоб иметь какое-то общее представление о порядке величин, и чтоб отслеживать регрессии в библиотеке, так что принимайте их соответственно; весь исходный код и данные есть в репозитории. Кроме того, нигде не буду писать об асимптотической сложности различных операций, т.к. не исследовал ее. В теории должно быть как у trie из википедии (например, получение элемента или поиск по префиксу — O(m), где m — длина ключа), но детали реализации (как libdatrie, так и моей обертки) могут все менять; если кто-нибудь построит соответствующие графики, буду, конечно же, благодарен.
Операции получения элемента, проверки in, обновления элемента работают в среднем раза в 2-3 медленнее, чем у стандартного словаря (т.е. быстро, на моем буке это от 1 до 3 млн операций в секунду для вышеупомянутого trie). Исключение — это вставка нового значения в trie, она работает значительно медленнее (порядка 50тыс операций в секунду для того же trie с русскими и английскими словами). При этом trie на таких данных занимает значительно меньше места в оперативной памяти: 3-5M (в зависимости от интерпретатора) vs 20M+ у обычного словаря (замеры памяти проводил коряво и за конкретные цифры не ручаюсь).
Т.е. datrie.Trie можно использовать как замену dict, если есть большое количество не очень длинных строк (например, слов или url'ов), данные в основном используются в режиме read-only (или «update-only») и хочется сэкономить оперативную память память ценой 2-3 кратного снижения скорости доступа.
Чтоб этот псевдо-read-only не очень смущал, trie можно сохранять в файл и загружать из файла:
>>> trie.save('my.trie')
>>> trie2 = datrie.Trie.load('my.trie')
Еще одна особенность trie — это то, что некоторые операции, которые в dict (и любой другой хэш-таблице) потребовали бы полного перебора или большого объема памяти для построения дополнительных индексов, в префиксном дереве работают почти так же быстро (а некоторые даже быстрее), как и простое получение значения.
Можно проверить, есть ли в trie элементы, ключи у которых начинаются с данного префикса (это даже быстрее, чем простая проверка in):
>>> trie.has_keys_with_prefix(u'fo')
True
Можно найти все префиксы данной строки, которые есть в данном trie (это медленнее, по тестам — 500-600тыс операций/сек):
>>> trie.prefixes(u'foobarbaz')
[u'foo', u'foobar']
>>> trie.prefix_items(u'foobarbaz')
[(u'foo', 5), (u'foobar', 10)]
Можно найти все элементы, у которых ключи начинаются с данной строки:
>>> trie.keys(u'fo')
[u'foo', u'foobar']
>>> trie.items(u'fo')
[(u'foo', 5), (u'foobar', 10)]
>>> trie.values(u'foob')
[10]
В последнем примере основное время сейчас тратится на построение результатов, а не на поиск, это можно оптимизировать; сейчас скорость была где-то порядка 150-300 тыс. значений/сек (например, при префиксе длиной 7 и в среднем 3 значениями это 70тыс операций/сек).
У datrie.Trie есть еще различные методы, справку по ним и дополнительные результаты замеров скорости можно посмотреть в README в репозитории.
Под pypy у меня все завелось на дебиане (раз в 10 медленнее, чем под cpython); на маке под pypy не завелось (с обычным питоном должно работать и на маке, и на линуксе, и под windows). C API — расширения в pypy всегда будут медленные, т.к. работают через костыль cpyext, а cython генерирует именно C API расширения. Можно было написать обертку на ctypes, но она была бы медленной под обычным питоном (и не факт что быстрой под pypy) + распространять ctypes-расширения неудобно. Ребята из pypy пилят сейчас cffi, обещают, что он быстрый будет (и под cpython, и под pypy). Плюс, возможно, cython научится когда-нибудь генерировать не C API расширения, а cffi-расширения. Вот тогда, возможно, и наступит счастье) А пока что с pypy делать — не знаю. Ничего не буду, наверное; как-то работает все под линуксом и ладно.
В процессе реализации столкнулся с тормознутым кодеком utf_32_le в питоне. На это есть баг с патчем ( bugs.python.org/issue15027 ), но патч еще не закоммитили. Изначально все операции в datrie работали в 10 раз медленнее, но потом удалось в одном месте обойтись без кодирования строки стандартным питоньим utf_32_le и все заработало получше. Этот кодек используется еще в паре «горячих» мест, так что, думаю, если его ускорят, некоторые операции в datrie могут заработать еще до 2 раз быстрее.
Итерация по дереву сейчас тоже не самая эффективная, это связано с особенностями интерфейса libdatrie. Но автор libdatrie — отличный чел и собирается все поправить, так что перспективы неплохие.
Как обычно, патчи, баг-репорты, идеи, бенчмарки, пулл-реквесты и тд приветствуются!
github / bitbucket, кому что удобнее.
Автор: kmike
