Basic Trie Implementation including Flow annotations

BasicTrie.js

// @flow

/**
 *
 * Trie Example
 * This is part of a series of examples that will lead to implemeting
 * a hashed mapped trie based immutable data structure.
 *
 */

// A TrieNode contains the value as well as information wether the node is
// completed and any further children nodes
class TrieNode {
  /*:: data: any*/
  /*:: isComplete: bool*/
  /*:: children: {[any] : TrieNode}*/
  constructor(data, isComplete = false) {
    // Data can be any value, i.e. a character
    this.data = data;
    // Defines if the node is completed, i.e. at the end.
    this.isComplete = isComplete;
    // A hash containing children nodes, could also implemented using a Map
    this.children = {};
  }

  // Enable to add and remove children TrieNodes

  /*:: addChild : (any, bool) => TrieNode*/
  addChild(data, isComplete = false) {
    // Check if we already have a reference to child node and if not
    // create a new TrieNode and add to children
    if (!this.hasChild(data)) {
      this.children[data] = new TrieNode(data, isComplete);
    }
    return this.children[data];
  }

  /*:: removeChild: any => bool*/
  removeChild(data /*:: : any */) /*:: : bool*/ {
    // Check if we have an existing child reference and remove if we have one
    if (this.hasChild(data)) {
      delete this.children[data];
      return true;
    }
    return false;
  }

  // Basic lookup either return a specific child or the complete hash.
  /*:: getChild: (any) => TrieNode */
  getChild(data) {
    if (data) {
      return this.children[data];
    }
  }

  // Check if we have a reference to the child
  /*:: hasChild: any => bool */
  hasChild(data) {
    return this.children[data] !== undefined;
  }

  // Returns all possible suggestions, i.e. ['a', 'e', 'i']
  /*:: getSuggestions : void => Array<string>*/
  getSuggestions() {
    return Object.keys(this.children);
  }

  // Basic size calculation
  /*:: childrenSize: void => number*/
  childrenSize() {
    return Object.keys(this.children).length;
  }

  // Basic check if node is completed
  /*:: isCompleted: void => bool*/
  isCompleted() {
    return this.isComplete;
  }

  // print data + show if completed + all possible suggestions, i.e.
  // "a:b,c" or "e (completed):f,g"
  /*:: toString: void => string*/
  toString() {
    return `${this.data}${
      this.isComplete ? " (completed)" : ""
    }:${this.getSuggestions().join(",")}`;
  }
}

// Tests
console.log("**Test Trie Node**");
const testTrieNode = new TrieNode("a", false);
testTrieNode.addChild("b", false);
testTrieNode.addChild("c", false);
console.log("Suggestions = ['b', 'c'], result:", testTrieNode.getSuggestions());
console.log("Contains 'b': ", testTrieNode.hasChild("b"));
console.log("toString = 'a:b,c', result:", testTrieNode.toString());

const testTrieNode2 = new TrieNode("e", true);
testTrieNode2.addChild("f", false);
testTrieNode2.addChild("g", false);
console.log(
  "Suggestions = ['f', 'g'], result:",
  testTrieNode2.getSuggestions()
);
console.log("Contains f: ", testTrieNode2.hasChild("f") === true);
console.log(
  "toString = 'e (completed): f,g', result:",
  testTrieNode2.toString()
);

console.log("-------------------------------");

// Constants
const EMPTY_STRING = "";

// A basic Trie Data Structure
class Trie {
  /*:: root: TrieNode*/
  constructor() {
    // We initialize with a root node containing an empty TrieNode
    // defined as an empty string
    this.root = new TrieNode(EMPTY_STRING);
  }

  // data could be a word i.e.
  /*:: add : (string | Array<string>) => void*/
  add(data) {
    this.addNode(this.root, data);
  }

  // We have to recurrsively go through the complete structure, starting
  // with the root node and then add the data, i.e. a word by breaking it a part
  // and adding nodes where they don't exist.
  /*:: addNode : (TrieNode, Array<string>) => void*/
  addNode(node, [key, ...rest]) {
    // We're at the end. We either have no node left or no more key to add.
    if (!node || !key) {
      return;
    }
    // Check if the node already has a reference to the child if not add
    // the data to the node's children. If we have no more data left
    // (via checking if we have any rest) we can also set the node as completed.
    if (!node.hasChild(key)) {
      node.addChild(key, rest.length === 0);
    }
    // Finally we call addNode with the child node and the rest of the data.
    this.addNode(node.getChild(key), rest);
  }

  // data could be a word i.e.
  /*:: remove : (string | Array<string>) => void */
  remove(data) {
    this.removeNode(this.root, data);
  }
  /*:: removeNode : (TrieNode, Array<string>) => void */
  removeNode(node, [key, ...rest]) {
    // We're at the end. We either have no node left or no more key to add.
    if (!node || !key) {
      return;
    }
    // We check if we have the key and do nothing if we don't have a key.
    if (node.hasChild(key)) {
      let child = node.getChild(key);
      // Do we have any other keys?
      if (rest.length > 0) {
        // If we only have one key left, we can remove the complete node
        if (child.childrenSize() === 0) {
          node.removeChild(key);
        } else {
          // Otherwise we continue with the child node for given key and the
          // rest of the data
          this.removeNode(child, rest);
        }
        // We have no more keys left
      } else {
        // check if we have any other child elements.
        // If we have no children left, we can remove the node!
        if (child.childrenSize() === 0) {
          node.removeChild(key);
        } else {
          // Otherwise we know that this child node is incomplete
          child.isComplete = false;
        }
      }
    }
  }

  // data can be a word i.e.
  /*:: contains : (string | Array<string>) => bool*/
  contains(data) {
    return this.containsNode(this.root, data);
  }

  /*:: containsNode : (TrieNode, Array<string>) => bool*/
  containsNode(node, [key, ...rest]) {
    // We're at the end. We either have no node left or no more key to add.
    if (!node || !key) {
      return false;
    }
    // If the node contains the key,
    if (node.hasChild(key)) {
      let child = node.getChild(key);
      // We have no data left and child node is completed. We found the node.
      if (child.isCompleted() && rest.length === 0) {
        return true;
      }
      // node is not completed. We need to call containsNode with the
      // child node and the rest of the data
      return this.containsNode(child, rest);
    }
    // Nothing found
    return false;
  }

  // Reset the root to an empty TrieNode.
  /*:: clear : void => void*/
  clear() {
    this.root = new TrieNode(EMPTY_STRING);
  }

  /*:: size : void => number*/
  size() {
    // We need to initialize the queue with the root node.

    let queue /*:: : Array<TrieNode>*/ = [this.root];
    let count = 0;
    while (queue.length) {
      // check if the current node is completed and if yes increment the count
      const node = queue.shift();
      if (node.isComplete) {
        count++;
      }
      // Retrieve the current node suggestions and add to the queue
      node.getSuggestions().forEach(key => {
        queue.push(node.getChild(key));
      });
    }
    // Queue is empty, we can return the actual count
    return count;
  }

  // Retrieve all completed words
  /*:: getCompleted : void => Array<string>*/
  getCompleted() {
    return this.getCompletedNode(this.root);
  }

  /*:: getCompletedNode : (TrieNode, ?Array<string>, ?string) => Array<string> */
  getCompletedNode(node, words = [], word = "") {
    // initialized with an empty words array and the root node
    // We iterate through all child nodes for a given node
    for (let key in node.children) {
      // Extend the current word with the next key
      const nextWord = word + key;
      // Check for every child node if completed and add to words if completed
      let child = node.getChild(key);
      if (child.isCompleted()) {
        words.push(nextWord);
      }
      // Call getCompletedNode with current child and the current words and word
      this.getCompletedNode(child, words, nextWord);
    }
    return words;
  }

  // Return any suggestions for a given data, i.e. "te" returns ["ten", "test"]
  /*:: suggestions : (string | Array<string>) => Array<string>*/
  getSuggestions(data) {
    return this.getSuggestionsByNode(this.root, data);
  }

  // Implementation is similar to getCompletedNode
  /*:: getSuggestionsByNode : (TrieNode, string, ?Array<string>, ?string) => Array<string> */
  getSuggestionsByNode(node, data, suggestions = [], word = "") {
    // initialized with an empty suggestions array and the root node
    // We iterate through all child nodes for a given node
    for (let key in node.children) {
      // Extend the current word with the next key
      const nextWord = word + key;
      // Check for every child node if completed
      // and if data overlaps the current word and add to suggestions
      let child = node.getChild(key);
      if (nextWord.indexOf(data) === 0 && child.isCompleted()) {
        suggestions.push(nextWord);
      }
      // Call getSuggestionsByNode with current child, the data for suggestions
      // and current suggestions array and word
      this.getSuggestionsByNode(child, data, suggestions, nextWord);
    }
    return suggestions;
  }
}

console.log("**Test Trie**");
const t = new Trie();
t.add("test");
t.add("ten");
t.add("fast");
t.add("slow");
console.log("Size === 4", t.size() === 4);
console.log("Contains test: ", t.contains("test") === true);
t.remove("test");
console.log("Does not contain test: ", t.contains("test") === false);
console.log("Size === 3", t.size() === 3);
console.log("Completed items = ['ten', 'fast', 'slow'] ", t.getCompleted());
t.add("technical");
t.add("testing");
console.log(
  "Suggestions for 'te' = ['testing', 'ten', 'technical'] ",
  t.getSuggestions("te")
);