Префиксное дерево с битмап-индексами

Не так давно мне потребовалось реализовать функционал, который уже был реализован другими людьми и не один раз, но не устраивал по каким то характеристикам. В данном случае требовалась какая-нибудь структура данных с возможностью поиска по строковому ключу или нескольким начальным символам ключа. Ключ — строка из латинских букв в любом регистре, пробелов и цифр.

Другие реализации не устроили, в первую очередь, по потребляемой памяти. Да и в любом случае ещё одна реализация префиксного дерева будет не лишней.

Идея была навеяна, видимо, этой статьёй, так как для индексирования в узле ссылок на последующие узлы было решено использовать битмап-индекс.

Очевидно, что число ссылок в одном узле может быть не более 63 (10 цифр, по 26 букв в каждом регистре плюс пробел). Поэтому для битмап-индекса в каждом узле используется 64-битное число. Установленный бит в какой-либо позиции означает наличие соответствующего символа и узла (поддерева).

В битмап-индексе распределение следующее: цифры от 0 до 9 занимают с 0 по 9 биты, буквы a-z с 10 по 35 биты, буквы A-Z с 36 по 62 биты и пробел занимает 63 бит. Для получения номера бита просто используется ASCII-код символа, уменьшенный на определённое число. Для каждого диапазона символов число своё: для цифр 48 (именно с 48 начинается диапазон цифр в ASCII-кодах), для букв a-z 87 (с 97 начинается диапазон букв a-z минус уже занятые цифрами 10 бит) и так далее.

Ниже показана соответствующая таблица





Узел представлен кодом ниже

type Trie struct {
    rw sync.RWMutex
    characters uint64
    subTries []*Trie
    data interface{}
}

Таким образом, каждый узел хранит битмап-индекс следующих символов, срез со ссылками на поддеревья и непосредственно данные. Расскажу кратко про механизмы поиска, вставки и удаления.

При поступлении на вход ключ посимвольно проходит через дерево, начиная с корневого узла. Для каждого символа из его ASCII-кода вычисляется номер бита, а на его основе вычисляется индекс в срезе ссылок (subTries) на поддеревья. Значение индекса есть число бит до искомого бита. Например, у нас есть такой индекс: 00...101. Это значит, что в битмап-индексе есть числа 0 и 2. Тогда индекс для числа 0 будет равен нулю, а для числа 2 единице, т. е. subTries будет хранить 2 ссылки на поддеревья: subTries[0] и subTries[1].

Код получения номера бита для символа:

func calcBitNum(char rune) (uint64, bool) {
	// Characters a-z use bit positions 10-35
	if char > 96 && char < 123 {
		return uint64(char - 87), true
	}

	// Characters A-Z use bit positions 36-61
	if char > 64 && char < 91 {
		return uint64(char - 29), true
	}

	// digits 0-9 use bit positions 0-9
	if char > 47 && char < 58 {
		return uint64(char - 48), true
	}

	// space uses 62 bit position
	if char == 32 {
		return 62, true
	}

	return 0, false
}

Код получения индекса узла для символа:

func (trie *Trie) getSubTreeIndex(bitNum uint64) int {
	shifted := trie.characters << (64 - bitNum)
	return bits.OnesCount64(shifted)
}

Код получения узла для заданного символа:

func (trie *Trie) getSubTrie(char rune) (uint64, int, *Trie) {
	bitNum, _ := calcBitNum(char)
	subTrieId := trie.getSubTreeIndex(bitNum)

	// There is no subTrie under given character since the corresponding bit is zero
	if trie.characters&(1<<bitNum) == 0 {
		return bitNum, subTrieId, nil
	}

	return bitNum, subTrieId, trie.subTries[subTrieId]
}

Вставка

Для вставки элемента выполняется проход по дереву. Если какой то символ (а значит, и все последующие) отсутствует, он вставляется в битмап-индекс. Создаётся поддерево, ссылка на которое вставляется в нужное место среза subTries (нужное место определяется числом бит до искомого). Когда дошли до последнего символа ключа, то в последнее поддерево вставляются данные.

Поиск

Для поиска ключ посимвольно «пробегает» по дереву. Как только он закончился, то возвращается результат. Если для ключа данные отсутствуют, но есть поддеревья далее, то выполняется выборка результатов из всех последующих поддеревьев, имеющих данные. Лимит на число возвращаемых данных можно задавать. Таким образом, можно делать поиск одного узла по совпадению ключа целиком, либо набора узлов по совпадению нескольких начальных символов.

Удаление

Процедура удаление вышла посложнее (что и ожидаемо), чем остальные. Для удаления ключа выполняется проход по дереву, в процессе которого в отдельный срез записываются узлы, не имеющие данных. Если встречается узел с данными, то срез обнуляется, а проход по дереву продолжается далее. То есть срез узлов на удаление должен содержать последовательность узлов, не имеющих данных. Если в итоге срез не пустой, то он проходится в обратном порядке и происходит удаление соответствующих бит из битмап-индекса и удаление узлов.

Итог

В итоге получилось префиксное дерево, превосходящее другие аналоги по скорости работы и потреблению памяти (в среднем на 30% меньше по памяти и на столько же быстрее). Также не поленился сделать возможным многопоточные операции с деревом, что существенно увеличивает производительность, особенно на операциях поиска. Само собой, решение довольно специализированное (ограничение в виде латинских символов, цифр и пробела в качестве составляющих ключа).

Исходя из назначения префиксного дерева данное решение можно использовать для создания словаря, поискового индекса.

Коды вставки, поиска и удаления приводить не стал, дабы не загромождать статью, к тому же, исходный код доступен на гитхабе.