Skip to content

Instantly share code, notes, and snippets.

@tpae
Created November 20, 2016 23:49
Show Gist options
  • Save tpae/72e1c54471e88b689f85ad2b3940a8f0 to your computer and use it in GitHub Desktop.
Save tpae/72e1c54471e88b689f85ad2b3940a8f0 to your computer and use it in GitHub Desktop.
Trie.js - super simple JavaScript implementation
// Trie.js - super simple JS implementation
// https://en.wikipedia.org/wiki/Trie
// -----------------------------------------
// we start with the TrieNode
function TrieNode(key) {
// the "key" value will be the character in sequence
this.key = key;
// 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;
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() {
this.root = new TrieNode(null);
}
// 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]);
// 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;
}
}
};
// 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' ]
@flopperj
Copy link

flopperj commented May 1, 2020

Awesome implementation of the Trie, here's my approach at implementing the remove method.

/**
 * Removes word from trie
 * Time complexity = O(nlogn)
 * Space complexity = O(n)
 * @param {string} word
 */
Trie.prototype.remove = function (word) {
    let root = this.root;

    if(!word) return;

    // recursively finds and removes a word
    const removeWord = (node, word) => {

        // check if current node contains the word
        if (node.end && node.getWord() === word) {

            // check and see if node has children
            let hasChildren = Object.keys(node.children).length > 0;

            // if has children we only want to un-flag the end node that marks end of a word.
            // this way we do not remove words that contain/include supplied word
            if (hasChildren) {
                node.end = false;
            } else {
                // remove word by getting parent and setting children to empty dictionary
                node.parent.children = {};
            }

            return true;
        }

        // recursively remove word from all children
        for (let key in node.children) {
            removeWord(node.children[key], word)
        }

        return false
    };

    // call remove word on root node
    removeWord(root, word);
};

@harshith-venkatesh
Copy link

This is amazing!

@by12380
Copy link

by12380 commented Sep 24, 2020

Thanks for the contribution! Just wanted to point out that unshift() in getWord() is not very efficient. unshift() is a O(n) operation and will lead to O(k^2) for getWord(). We would rather do push() and then reverse() in the end.

@AndrewKnutsen
Copy link

I'm too new to JS not to ask: does this need a destroy() method of some sort?

@docdawning
Copy link

I agree with @Cxarli, an explicit license would be nice.

@lancejpollard
Copy link

class TrieNode {
  constructor(key) {
    // the "key" value will be the character in sequence
    this.key = key;

    // 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;
  }

  getWord() {
    let output = [];
    let node = this;

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

    return output.join('')
  }
}

class Trie {
  constructor() {
    this.base = new TrieNode(null)
  }

  insert(word) {
    let node = this.base

    const points = Array.from(word)

    for (const i in points) {
      const point = points[i]
      if (!node.children[point]) {
        node.children[point] = new TrieNode(point)
        node.children[point].parent = node
      }

      node = node.children[point]

      if (i == word.length - 1) {
        node.end = true
      }
    }
  }

  contains(word) {
    let node = this.base

    const points = Array.from(word)

    for (const i in points) {
      const point = points[i]

      if (node.children[point]) {
        node = node.children[point]
      } else {
        return false
      }
    }

    return node.end;
  }

  find(prefix) {
    let node = this.base
    let output = []

    const points = Array.from(prefix)

    for (const i in points) {
      const point = points[i]

      // make sure prefix actually has words
      if (node.children[point]) {
        node = node.children[point]
      } else {
        // there's none. just return it.
        return output
      }
    }

    const stack = [node]
    while (stack.length) {
      node = stack.shift()
      // base case, if node is at a word, push to output
      if (node.end) {
        output.unshift(node.getWord())
      }

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

    return output
  }
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment