Trie.ts - super simple JavaScript implementation

Trie.ts

// Trie.ts - super simple JS implementation
// https://en.wikipedia.org/wiki/Trie

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

// we start with the TrieNode
class TrieNode<V> {
  // we keep a reference to parent
  public parent = null;
  
  // we have hash of children
  public children = {};
  
  // check to see if the node is at the end
  public end = false;

  // the "key" value will be the character in sequence
  constructor(public key: string, public value?: V) {}

  // iterates through the parents to get the word.
  // time complexity: O(k), k = word length
  getWord() {
    var output = [];
    var node = this;
    
    while (node !== null) {
      output.unshift(node.key);
      node = node.parent;
    }
    
    return output.join('');
  };
}


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

// we wrap our find result in class to include its value
class TrieFindResult<V> {
  constructor(public word: string, public value?: V) {}
}


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

// we implement Trie with just a simple root with null value.
class Trie<V> {
  public root = new TrieNode<V>(null);

  // inserts a word into the trie.
  // time complexity: O(k), k = word length
  insert(word: string, val?: V) {
    var node = this.root; // we start at the root 😬
    
    // for every character in the word
    for(var i = 0; i < word.length; i++) {
      // check to see if character node exists in children.
      if (!node.children[word[i]]) {
        // if it doesn't exist, we then create it.
        node.children[word[i]] = new TrieNode<V>(word[i]);
        
        // we also assign the parent to the child node.
        node.children[word[i]].parent = node;
      }
      
      // proceed to the next depth in the trie.
      node = node.children[word[i]];
      
      // finally, we check to see if it's the last word.
      if (i == word.length-1) {
        // if it is, we set the end flag to true.
        node.end = true;
        node.value = val;
      }
    }
  };

  // check if it contains a whole word.
  // time complexity: O(k), k = word length
  contains(word: string) {
    var node = this.root;
    
    // for every character in the word
    for(var i = 0; i < word.length; i++) {
      // check to see if character node exists in children.
      if (node.children[word[i]]) {
        // if it exists, proceed to the next depth of the trie.
        node = node.children[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
  find(prefix: string): TrieFindResult<V>[] {
    var node = this.root;
    var output: TrieFindResult<V>[] = [];
    
    // for every character in the prefix
    for(var i = 0; i < prefix.length; i++) {
      // make sure prefix actually has words
      if (node.children[prefix[i]]) {
        node = node.children[prefix[i]];
      } else {
        // there's none. just return it.
        return output;
      }
    }
    
    // recursively find all words in the node
    this.findAllWords(node, output);
    
    return output;
  }

  // recursive function to find all words in the given node.
  private findAllWords(node: TrieNode<V>, arr: TrieFindResult<V>[]) {
    // base case, if node is at a word, push to output
    if (node.end) {
      arr.unshift(new TrieFindResult(node.getWord(), node.value));
    }
    
    // iterate through each children, call recursive findAllWords
    for (var child in node.children) {
      this.findAllWords(node.children[child], arr);
    }
  }
}

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

// instantiate our trie
var trie = new Trie();

// insert few values
trie.insert("hello", 3);
trie.insert("helium", 5);

// check contains method
console.log(trie.contains("helium"));  // true
console.log(trie.contains("kickass")); // false

// check find method
console.log(trie.find("hel"));  // [ { word: 'helium', value: 5 }, { word: 'hello', value: 3 } ]
console.log(trie.find("hell")); // [ { word: 'hello', value: 3 } ]