JavaScript/Data-Structures/Tree/Trie.js

class TrieNode {
  constructor(key, parent) {
    this.key = key
    this.count = 0
    this.children = Object.create(null)
    if (parent === undefined) {
      this.parent = null
    } else {
      this.parent = parent
    }
  }
}

class Trie {
  constructor() {
    // create only root with null key and parent
    this.root = new TrieNode(null, null)
  }

  // Recursively finds the occurrence of all words in a given node
  static findAllWords(root, word, output) {
    if (root === null) return
    if (root.count > 0) {
      if (typeof output === 'object') {
        output.push({ word, count: root.count })
      }
    }
    let key
    for (key in root.children) {
      word += key
      this.findAllWords(root.children[key], word, output)
      word = word.slice(0, -1)
    }
  }

  insert(word) {
    if (typeof word !== 'string') return
    if (word === '') {
      this.root.count += 1
      return
    }
    let node = this.root
    const len = word.length
    let i
    for (i = 0; i < len; i++) {
      if (node.children[word.charAt(i)] === undefined) {
        node.children[word.charAt(i)] = new TrieNode(word.charAt(i), node)
      }
      node = node.children[word.charAt(i)]
    }
    node.count += 1
  }

  findPrefix(word) {
    if (typeof word !== 'string') return null
    let node = this.root
    const len = word.length
    let i
    // After end of this loop node will be at desired prefix
    for (i = 0; i < len; i++) {
      if (node.children[word.charAt(i)] === undefined) return null // No such prefix exists
      node = node.children[word.charAt(i)]
    }
    return node
  }

  remove(word, count) {
    if (typeof word !== 'string') return
    if (typeof count !== 'number') count = 1
    else if (count <= 0) return

    // for empty string just delete count of root
    if (word === '') {
      if (this.root.count >= count) this.root.count -= count
      else this.root.count = 0
      return
    }

    let child = this.root
    const len = word.length
    let i, key
    // child: node which is to be deleted
    for (i = 0; i < len; i++) {
      key = word.charAt(i)
      if (child.children[key] === undefined) return
      child = child.children[key]
    }

    // Delete no of occurrences specified
    if (child.count >= count) child.count -= count
    else child.count = 0

    // If some occurrences are left we don't delete it or else
    // if the object forms some other objects prefix we don't delete it
    // For checking an empty object
    // https://stackoverflow.com/questions/679915/how-do-i-test-for-an-empty-javascript-object
    if (
      child.count <= 0 &&
      Object.keys(child.children).length &&
      child.children.constructor === Object
    ) {
      child.parent.children[child.key] = undefined
    }
  }

  findAllWords(prefix) {
    const output = []
    // find the node with provided prefix
    const node = this.findPrefix(prefix)
    // No such prefix exists
    if (node === null) return output
    Trie.findAllWords(node, prefix, output)
    return output
  }

  contains(word) {
    // find the node with given prefix
    const node = this.findPrefix(word)
    // No such word exists
    return node !== null && node.count !== 0
  }

  findOccurrences(word) {
    // find the node with given prefix
    const node = this.findPrefix(word)
    // No such word exists
    if (node === null) return 0
    return node.count
  }
}

export { Trie }