algorithmpool/src/dspool/tree/trie.js

A pool of algorithms and data-structures for geeks

const validate = (str) => {
    if (typeof str !== 'string') {
        throw new Error(`${str} is not a string`);
    }
};

export const node = (ch) => {
    const char = ch;
    let endOfWord = false;
    let parent = null;
    const children = {};
    let childrenCount = 0;

    validate(ch);

    const getChar = () => char;

    const setParent = (p) => {
        parent = p;
    };

    const getParent = () => parent;

    const setEndOfWord = (eow) => {
        endOfWord = eow;
    };

    const isEndOfWord = () => endOfWord;

    const addChild = (child) => {
        children[child.getChar()] = child;
        childrenCount += 1;
    };

    const removeChild = (child) => {
        children[child.getChar()] = null;
        childrenCount -= 1;
    };

    const getChild = c => children[c] || null;

    const getChildren = () => children;

    const countChildren = () => childrenCount;

    // trie node api
    return {
        getChar,
        setParent,
        getParent,
        setEndOfWord,
        isEndOfWord,
        addChild,
        removeChild,
        getChild,
        getChildren,
        countChildren
    };
};
export class trie {

    constructor(){
        this.nodesCount = 1;
        this.wordsCount = 0;
        this.rootNode = node('');
    }

    node(ch){
        return node(ch);
    }
    
    countNodes(){
        this.nodesCount;
    }

    countWords(){
        this.wordsCount;
    }

    search(word){
        validate(word);
        let currentNode = this.rootNode;
        for (let i = 0; i < word.length; i += 1) {
            const child = currentNode.getChild(word[i]);
            if (child === null) {
                return null;
            }
            currentNode = child;
        }
        if (currentNode.isEndOfWord()) {
            return currentNode;
        }
        return null;
    };

    insert(word){
        validate(word);
        let currentNode = this.rootNode;
        for (let i = 0; i < word.length; i += 1) {
            if (currentNode.getChild(word[i]) === null) {
                const child = node(word[i]);
                child.setParent(currentNode);
                currentNode.addChild(child);
                this.nodesCount += 1;
            }
            currentNode = currentNode.getChild(word[i]);
        }
        if (currentNode.getChar() !== '' && !currentNode.isEndOfWord()) {
            currentNode.setEndOfWord(true);
            this.wordsCount += 1;
        }
    };

    remove(word){
        let currentNode = this.search(word);
        if (currentNode !== null && currentNode.getChar() !== '') {
            if (currentNode.countChildren() > 0) {
                currentNode.setEndOfWord(false);
            } else {
                while (currentNode.getParent() !== null) {
                    if (currentNode.countChildren() === 0) {
                        currentNode.getParent().removeChild(currentNode);
                        this.nodesCount -= 1;
                    }
                    currentNode = currentNode.getParent();
                }
            }
            this.wordsCount -= 1;
        }
    };

    traverse(cb){
        let word = '';
        const traverseFn = (currentNode) => {
            if (currentNode.isEndOfWord()) {
                cb(word);
            }
            const children = currentNode.getChildren();
            Object.keys(children).forEach((char) => {
                const child = children[char];
                word += child.getChar();
                traverseFn(child); // depth-first traverse
                word = word.substr(0, word.length - 1); // word backward tracking
            });
        };
        return traverseFn(this.rootNode);
    };

    clear(){
        this.nodesCount = 1;
        this.wordsCount = 0;
        this.rootNode = node('');
    };

}