@picosearch/radix-tree
Version:
Simple, zero dependency, type-safe implementation of a radix tree data structure.
230 lines (229 loc) • 8.98 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.RadixTree = exports.RadixTreeMap = void 0;
const constants_1 = require("./constants");
const util_1 = require("./util");
class RadixTreeMap {
root;
constructor(root = (0, util_1.getNewEmptyNode)()) {
this.root = root;
}
/**
* Inserts a key into the tree.
*
* @param key The key to insert.
* @param values The values to associate with the key.
*/
insert(key, ...values) {
(0, util_1.assert)(!!key && key.length > 0, 'Missing key!');
let keyOffset = 0;
let currentNode = this.root;
while (currentNode) {
const children = currentNode[0];
currentNode = null;
for (const [edgeLabel, childNode] of Object.entries(children)) {
const commonPrefix = (0, util_1.getCommonPrefix)(edgeLabel, key.slice(keyOffset));
if (!commonPrefix)
continue;
keyOffset += commonPrefix.length;
(0, util_1.assert)(keyOffset <= key.length, 'Unexpected key offset overflow');
(0, util_1.assert)(commonPrefix.length <= edgeLabel.length, 'Unexpected prefix length overflow');
if (commonPrefix.length === edgeLabel.length) {
if (keyOffset === key.length) {
childNode[1] ??= [];
childNode[1].push(...values);
return;
}
currentNode = childNode;
break;
}
// partial match => split the edge
children[commonPrefix] = (0, util_1.getNewEmptyNode)();
const suffixOld = edgeLabel.slice(commonPrefix.length);
children[commonPrefix][0][suffixOld] = childNode;
delete children[edgeLabel];
const suffixNew = key.slice(keyOffset);
if (suffixNew.length) {
children[commonPrefix][0][suffixNew] = (0, util_1.getNewEmptyNode)();
children[commonPrefix][0][suffixNew][1] ??= [];
children[commonPrefix][0][suffixNew][1].push(...values);
}
else {
children[commonPrefix][1] ??= [];
children[commonPrefix][1].push(...values);
}
return;
}
if (!currentNode) {
// none of the children match
const suffix = key.slice(keyOffset);
children[suffix] = (0, util_1.getNewEmptyNode)();
children[suffix][1] ??= [];
children[suffix][1].push(...values);
}
}
}
/**
* Returns the values for the given key.
*
* @param key The key to look up.
* @returns The values for the given key or null if the key is not found.
*/
lookup(key) {
if (!key?.length)
return null;
let keyOffset = 0;
let currentNode = this.root;
while (currentNode) {
const children = currentNode[0];
currentNode = null;
for (const [edgeLabel, childNode] of Object.entries(children)) {
const commonPrefix = (0, util_1.getCommonPrefix)(edgeLabel, key.slice(keyOffset));
if (!commonPrefix)
continue;
(0, util_1.assert)(commonPrefix.length <= edgeLabel.length, 'Unexpected prefix length overflow');
if (commonPrefix.length === edgeLabel.length) {
keyOffset += commonPrefix.length;
(0, util_1.assert)(keyOffset <= key.length, 'Unexpected key offset overflow');
if (keyOffset === key.length)
return childNode[1] ?? null;
currentNode = childNode;
break;
}
// the key matches only a part of the edge
return null;
}
}
return null;
}
getFuzzyMatches(query, options) {
if (!query?.length)
return [];
const maxErrors = options?.maxErrors ?? constants_1.FUZZY_SEARCH_DEFAULT_MAX_ERROR;
if (maxErrors === 0) {
// this case equals a regular non-fuzzy search...
const values = this.lookup(query);
if (values === null)
return [];
return [[query, values]];
}
const result = [];
// The following traverses the RadixTreeMap while computing the Edit distance for each path
// using dynamic programming. In this matrix the query characters are column indices and
// candidate words take the vertical axis. Only the last and the current row need to be
// remembered for each pair of words.
//
// see https://repositorio.uchile.cl/bitstream/handle/2250/126168/Navarro_Gonzalo_Guided_tour.pdf (page 17)
const stack = [
{
prefix: '',
// the first row in the matrix is just the sequence 0,1,..,sequence.length
prevRow: Array.from({ length: query.length + 1 }, (_, i) => i),
node: this.root,
},
];
traversal: while (true) {
const currentItem = stack.shift();
if (!currentItem)
break;
const { prefix, prevRow, node } = currentItem;
for (const item of Object.entries(node[0])) {
const [edgeLabel, childNode] = item;
let lastRow = prevRow;
let newPrefix = prefix;
let lastRowErrorsTooHigh = false;
for (const char of edgeLabel.split('')) {
newPrefix += char;
const newRow = [newPrefix.length];
for (let j = 1; j <= query.length; j++) {
const charDist = query[j - 1] === char ? 0 : 1;
const value = Math.min(lastRow[j - 1] + charDist, lastRow[j] + 1, newRow[j - 1] + 1);
newRow.push(value);
}
lastRow = newRow;
lastRowErrorsTooHigh = lastRow.every((v) => v > maxErrors);
if (lastRowErrorsTooHigh)
break;
}
const childNodeValues = childNode?.[1];
if (lastRow[lastRow.length - 1] <= maxErrors &&
childNode &&
childNodeValues &&
childNodeValues.length > 0) {
(0, util_1.sortedInsert)(result, {
match: newPrefix,
distance: lastRow[lastRow.length - 1],
values: childNodeValues,
}, (a, b) => {
const diff = a.distance - b.distance;
return diff === 0 ? a.match.localeCompare(b.match) : diff;
});
if (options?.limit && result.length >= options.limit)
break traversal;
}
if (!lastRowErrorsTooHigh) {
stack.push({
prefix: newPrefix,
prevRow: lastRow,
node: childNode,
});
}
}
}
return options?.includeValues
? result.map((item) => [item.match, item.values])
: result.map((item) => item.match);
}
toJSON() {
return JSON.stringify(this.root);
}
static fromJSON(jsonStr) {
return new RadixTreeMap(JSON.parse(jsonStr));
}
}
exports.RadixTreeMap = RadixTreeMap;
class RadixTree {
tree;
constructor(options) {
this.tree = new RadixTreeMap(options?.root ?? (0, util_1.getNewEmptyNode)());
}
/**
* Inserts a key into the tree.
*
* @param key The key to insert.
*/
insert(key) {
this.tree.insert(key, 1);
}
/**
* Checks if the tree contains the given key.
*
* @param key The key to check.
* @returns Whether the tree contains the given key.
*/
has(key) {
return this.tree.lookup(key) !== null;
}
/**
* Returns the fuzzy matches for the given query.
*
* @param query The query to search for.
* @param options The options to use for the search.
* @returns The fuzzy matches for the given query.
*/
getFuzzyMatches(query, options) {
return this.tree.getFuzzyMatches(query, {
includeValues: false,
...options,
});
}
toJSON() {
return this.tree.toJSON();
}
static fromJSON(jsonStr) {
return new RadixTree({
root: JSON.parse(jsonStr),
});
}
}
exports.RadixTree = RadixTree;