atomic-fns/dist/collections/Trie.js

Like Lodash, but for ESNext and with types. Stop shipping code built for browsers from 2015.

import { Collection } from './abc.js';
function* traverse(node, path, total = 20, count = 0) {
    if (node.isWord) {
        count++;
        yield path;
    }
    for (const key of Object.keys(node)) {
        yield* traverse(node[key], path + key, total, count);
    }
}
/**
 * A Trie is a tree for efficient storing and retrieval of strings such as words in a dictionary. They are useful for checking all the strings that have a common prefix.
 * Insertions, lookups and removals in `O(s)` where `s` is the search term.
 * @see {@link https://en.wikipedia.org/wiki/Trie Trie}
 */
export class Trie extends Collection {
    root = {};
    count = 0;
    constructor(words) {
        super();
        if (words) {
            for (const word of words) {
                this.add(word);
            }
        }
    }
    /**
     * Adds the given string to the tree.
     * @param word A string
     */
    add(word) {
        let current = this.root;
        for (let i = 0; i < word.length; i++) {
            const c = word[i];
            if (!current[c])
                current[c] = {};
            current = current[c];
        }
        // Set as end of the word
        Object.defineProperty(current, 'isWord', {
            value: true,
            enumerable: false,
            writable: true
        });
        this.count++;
        return this;
    }
    /**
     * Removes the given string from the tree.
     * @param word The string to remove
     * @returns {boolean} Returns `true` if the string was found and removed.
     */
    remove(word) {
        for (let i = 0, node = this.root; i < word.length && node; i++) {
            node = node[word[i]];
            if (i === word.length - 1 && node.isWord) {
                node.isWord = false;
                this.count--;
                return true;
            }
        }
        return false;
    }
    findNode(word) {
        let node = this.root;
        for (let i = 0; i < word.length && node; i++) {
            const letter = word[i];
            node = node[letter];
        }
        return node;
    }
    /**
     * Generates all results that start with the given string.
     * @param word A prefix string
     * @param {number} [limit=20] The number of results to return.
     */
    matches = function* (word, limit = 20) {
        const node = this.findNode(word);
        if (node) {
            yield* traverse(node, word, limit);
        }
    };
    /**
     * Returns `true` if the given string is found in the tree.
     * @param {string} word
     */
    contains(word) {
        const node = this.findNode(word);
        return node?.isWord;
    }
    /**
     * Removes all strings from the tree.
     */
    clear() {
        this.root = {};
        this.count = 0;
    }
    /**
     * Returns the total number of strings in the tree.
     */
    get size() {
        return this.count;
    }
}