Trie.js - super simple JavaScript implementation

Trie.js

// Trie.js - super simple JS implementation
// https://en.wikipedia.org/wiki/Trie
// From https://gist.github.com/tpae/72e1c54471e88b689f85ad2b3940a8f0

// Modified to support key -> value mapping for quick object retrieval,
// using an ES6 Map instead of plain JS object, and optional case insensitivity.
// Values are automatically converting to arrays to support multiple values
// mapped to a single tree leaf.

//  Example:
//  
//  let trie = new Trie();
//  
//  // support for multiple isotopes (aka one-to-many values per leaf node)
//
//  trie.insert("helium", { protons : 2, neutrons : 2 }); // He-3
//  trie.insert("helium", { protons : neutrons : 3 }); // He-4
//  trie.insert("lithium", { protons : 3, neutrons : 3 }); // Li-6
//  trie.insert("lithium", { protons : 3, neutrons : 4 }); // Li-7
//  
//  console.log(trie.contains("lit")); // true
//  console.log(trie.find("hel", true)); // [ protons : 2, neutrons : 2, { protons : neutrons : 3 } ]
// 

// -----------------------------------------

(function(window) {

    // recursive function to find all words in the given node.
    const findAllWords = function(node, arr, returnValue) {
        // base case, if node is at a word, push to output
        if (node.end) {
            if (returnValue) {
                if (Array.isArray(node.value)) {
                    Array.prototype.unshift.apply(arr, node.value);
                } else {
                    arr.unshift(node.value);
                }
            } else {
                arr.unshift(node.getWord());
            }
        }
        
        // iterate through each children, call recursive findAllWords
        for (let [child] of node.children.entries()) {
            findAllWords(node.children.get(child), arr, returnValue);
        }
    };

    // we start with the TrieNode
    function TrieNode(key, value) {
        // the "key" value will be the character in sequence
        this.key = key;
        this.value = value;
        
        // we keep a reference to parent
        this.parent = null;
        
        // we have hash of children
        this.children = new Map();
        
        // check to see if the node is at the end
        this.end = false;
    }
    
    // iterates through the parents to get the word.
    // time complexity: O(k), k = word length
    TrieNode.prototype.getWord = function() {
        let output = [];
        let node = this;
        
        while (node !== null) {
            output.unshift(node.key);
            node = node.parent;
        }
        
        return output.join('');
    };
    
    // -----------------------------------------
    
    // we implement Trie with just a simple root with null value.
    function Trie (options) {
        this.root = new TrieNode(null);
        this.options = options || Object.create(null);
    }
    
    // inserts a word into the trie.
    // time complexity: O(k), k = word length
    Trie.prototype.insert = function(inWord, value) {
        let word = inWord;
        if (!(word && word.length)) return;

        let node = this.root; // we start at the root 😬
        if (this.options.case_insensitive) {
            word = word.toLowerCase();
        }
        
        // for every character in the word
        for(let i = 0, e = word.length; i < e; i++) {
            // check to see if character node exists in children.
            if (!node.children.has(word[i])) {
                // if it doesn't exist, we then create it.
                let iNode = new TrieNode(word[i]);
                
                // we also assign the parent to the child node.
                iNode.parent = node;
                node.children.set(word[i], iNode);
            }
            
            // proceed to the next depth in the trie.
            node = node.children.get(word[i]);
            
            // finally, we check to see if it's the last word.
            if (i === e - 1) {
                // if it is, we set the end flag to true.
                node.end = true;

                if (node.value) {
                    node.value.push( value );
                } else {
                    node.value = [ value ];
                }
            }
        }
    };
    
    // check if it contains a whole word.
    // time complexity: O(k), k = word length
    Trie.prototype.contains = function(inWord) {
        let word = inWord;
        if (!(word && word.length)) return false;
        let node = this.root;

        if (this.options.case_insensitive) {
            word = word.toLowerCase();
        }
        
        // for every character in the word
        for (let i = 0, e = word.length; i < e; i++) {
            // check to see if character node exists in children.
            if (node.children.has(word[i])) {
                // if it exists, proceed to the next depth of the trie.
                node = node.children.get(word[i]);
            } else {
                // doesn't exist, return false since it's not a valid word.
                return false;
            }
        }
        
        // we finished going through all the words, but is it a whole word?
        return node.end;
    };
    
    // returns every word with given prefix
    // time complexity: O(p + n), p = prefix length, n = number of child paths
    Trie.prototype.find = function(inPrefix, returnValue) {
        let prefix = inPrefix;
        let output = [];
        if (!(prefix && prefix.length)) return output;
        let node = this.root;

        if (this.options.case_insensitive) {
            prefix = prefix.toLowerCase();
        }
        
        // for every character in the prefix
        for (let i = 0, e = prefix.length; i < e; i++) {
            // make sure prefix actually has words
            if (node.children.has(prefix[i])) {
                node = node.children.get(prefix[i]);
            } else {
                // there's none. just return it.
                return output;
            }
        }
        
        // recursively find all words in the node
        findAllWords(node, output, returnValue);
        
        return output;
    };

    window.Trie = Trie;

})(window);