Trie.js - simple Trie implementation. Based off Wengrow's "A Common-Sense Guide to Data Structures and Algorithms in Python" (Chapter 17).

Trie.js

// Trie.js - simple Trie implementation
// Based off Wengrow's A Common-Sense Guide to Data Structures and Algorithms in Python (Chapter 17).

class TrieNode {
  constructor() {
    this.children = {};
  }
}

class Trie {
  constructor() {
    this.root = new TrieNode();
    this._WORD_END_MARKER = "*";
  }

  /*
  Insert a word in the trie. 
  */
  insert(word) {
    let currNode = this.root;

    for (let char of word) {
      if (!currNode.children[char]) {
        currNode.children[char] = new TrieNode(char);
      }

      currNode = currNode.children[char];
    }

    currNode.children[this._WORD_END_MARKER] = null;
  }

  /*
  Check if it contains whole word.
  */
  contains(word) {
    let currNode = this.root;

    for (const char of word) {
      if (!currNode.children[char]) return false

      currNode = currNode.children[char];
    }

    return currNode.children[this._WORD_END_MARKER] === null;
  }

  /*
  Return an array of all the words in the trie, starting from the root.
  If a node is provided, return all the words that start from that node.
  */
  getAllWords(node = null) {
    let words = [];
    const WORD_END_MARKER = this._WORD_END_MARKER;

    function collectAllWordsRec(words, node, word = "") {
      let currNode = node;
      for (const char of Object.keys(currNode.children)) {
        if (char === WORD_END_MARKER) {
          words.push(word);
        } else {
          const childNode = currNode.children[char];
          collectAllWordsRec(words, childNode, word + char);
        }
      }

      return words;
    }

    collectAllWordsRec(words, node || this.root);
    return words;
  }

  /*
  Return a TrieNode containing the last character of a word in the Trie. If
  a word is prefix, return the final character in the prefix.
  */
  search(word) {
    let currNode = this.root;
    for (const char of word) {
      if (char in currNode.children) {
        currNode = currNode.children[char];
      } else {
        return null;
      }
    }

    return currNode;
  }

  /*
  Return the strings that complete a prefix.
  */
  autocomplete(prefix) {
    let currNode = this.search(prefix);
    if (!currNode) {
      return null;
    } else {
      return this.getAllWords(currNode);
    }
  }

  /*
  Given a input word not in the trie, return the first word that shares the
  longest possible prefix with the input word. If the input word is in the
  trie, return it.
  */
  autocorrect(word) {
    let word_found_so_far = "";

    let currNode = this.root;
    for (const char of word) {
      if (currNode.children[char]) {
        word_found_so_far += char;
        currNode = currNode.children[char];
      } else {
        const currNodeSuffixes = this.getAllWords(currNode);
        return word_found_so_far + currNodeSuffixes[0];
      }
    }

    return word;
  }
}

// instantiate a trie.
const trie = new Trie();

// insert a few words.
trie.insert("computer");
trie.insert("computation");
trie.insert("communication");
trie.insert("command");
trie.insert("comedy");
trie.insert("colosseum");

// check contains method
console.log(trie.contains("computer")); // true
console.log(trie.contains("comp")); // false
console.log(trie.contains("laptop")); // false

// check autocomplete method.
console.log(trie.autocomplete("col")); // [ 'osseum' ]
console.log(trie.autocomplete("comm")); // [ 'unication', 'and' ]
console.log(trie.autocomplete("com")); // [ 'puter', 'putation', 'munication', 'mand', 'edy' ]

// check autocorrect method.
console.log(trie.autocorrect("computer")); // 'computer'
console.log(trie.autocorrect("compilation")); // 'computer'

// check search method.
console.log(trie.search("comed")); // TrieNode { children: { y: TrieNode { children: [Object] } } }
console.log(trie.search("colosseum")); // TrieNode { children: { '*': null } }

// check getAllWords method.
console.log(trie.getAllWords()); // [ 'computer', 'computation', ... ]
console.log(trie.getAllWords(trie.search("col"))); // [ 'osseum' ]