Trie.js - super simple JavaScript implementation

Trie.js

/* tslint:disable */

import { AutocompleteCandidate } from '../autocomplete/candidate';

// Source code token from below link
// https://gist.github.com/tpae/72e1c54471e88b689f85ad2b3940a8f0
// @MODIFIED has been added where modifications have been made
//
// Trie.js - super simple JS implementation
// https://en.wikipedia.org/wiki/Trie
//
// we start with the TrieNode
function TrieNode(key, confidence) {
  // the "key" value will be the character in sequence
  this.key = key;
  this.confidence = confidence;

  // we keep a reference to parent
  this.parent = null;

  // we have hash of children
  this.children = {};

  // 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() {
  var output = [];
  var node = this;
  var confidence = node.confidence; // @MODIFIED keep track of confidence for the word

  while (node !== null) {
    output.unshift(node.key);
    node = node.parent;
  }

  return new AutocompleteCandidate(output.join(''), confidence); // @MODIFIED return an object with word and confidence info
};

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

// we implement Trie with just a simple root with null value.
// @MODIFIED: Added export default
export default function Trie() {
  this.root = new TrieNode(null, 0);
  this.size = 0; // @MODIFIED keep track of the size of the trie
}

// inserts a word into the trie.
// time complexity: O(k), k = word length
Trie.prototype.insert = function(word) {
  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(word[i], 0); // @MODIFIED added initial confidence of 0 (this is just the trie node getting setup)
      this.size++; // @MODIFIED keep track of total number of nodes
      // 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.confidence++; // @MODIFIED add confidence increment
    }
  }
};

// check if it contains a whole word.
// time complexity: O(k), k = word length
Trie.prototype.contains = function(word) {
  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
Trie.prototype.find = function(prefix) {
  var node = this.root;
  var output = [];

  // 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
  findAllWords(node, output);

  return output;
};

// recursive function to find all words in the given node.
function findAllWords(node, arr) {
  // base case, if node is at a word, push to output
  if (node.end) {
    arr.unshift(node.getWord());
  }

  // iterate through each children, call recursive findAllWords
  for (var child in node.children) {
    findAllWords(node.children[child], arr);
  }
}

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

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

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

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

// check find method
console.log(trie.find("hel"));  // [ 'helium', 'hello' ]
console.log(trie.find("hell")); // [ 'hello' ]