Trie Autocomplete With Weighted Words

Store weighted words in a case-insensitive trie and return autocomplete suggestions ranked by weight and prefix match.

JavascriptCli Tooling

Code

snippet
class TrieNode {
  constructor() {
    this.children = new Map();
    this.terminal = false;
    this.weight = 0;
  }
}

export class Trie {
  #root = new TrieNode();
  #caseSensitive;
  #wordCount = 0;

  constructor({ caseSensitive = false } = {}) {
    if (typeof caseSensitive !== "boolean") {
      throw new TypeError("caseSensitive must be a boolean");
    }
    this.#caseSensitive = caseSensitive;
  }

  #normalize(word) {
    if (typeof word !== "string" || word.trim().length === 0) {
      throw new TypeError("word must be a non-empty string");
    }
    return this.#caseSensitive ? word : word.toLowerCase();
  }

  insert(word, weight = 1) {
    if (!Number.isFinite(weight) || weight < 0) {
      throw new RangeError("weight must be a non-negative finite number");
    }
    const normalized = this.#normalize(word);
    let node = this.#root;
    for (const character of normalized) {
      let child = node.children.get(character);
      if (child === undefined) {
        child = new TrieNode();
        node.children.set(character, child);
      }
      node = child;
    }
    if (!node.terminal) {
      node.terminal = true;
      this.#wordCount += 1;
    }
    node.weight = weight;
    return this;
  }

  has(word) {
    const normalized = this.#normalize(word);
    let node = this.#root;
    for (const character of normalized) {
      const child = node.children.get(character);
      if (child === undefined) return false;
      node = child;
    }
    return node.terminal;
  }

  startsWith(prefix) {
    if (typeof prefix !== "string") {
      throw new TypeError("prefix must be a string");
    }
    const normalized = this.#caseSensitive ? prefix : prefix.toLowerCase();
    let node = this.#root;
    for (const character of normalized) {
      const child = node.children.get(character);
      if (child === undefined) return false;
      node = child;
    }
    return true;
  }

  suggest(prefix, limit = 5) {
    if (typeof prefix !== "string") {
      throw new TypeError("prefix must be a string");
    }
    if (!Number.isInteger(limit) || limit < 1) {
      throw new RangeError("limit must be a positive integer");
    }
    const normalized = this.#caseSensitive ? prefix : prefix.toLowerCase();
    let node = this.#root;
    for (const character of normalized) {
      const child = node.children.get(character);
      if (child === undefined) return [];
      node = child;
    }

    const matches = [];
    this.#collect(node, normalized, matches);
    matches.sort(
      (left, right) =>
        right.weight - left.weight || left.word.localeCompare(right.word),
    );
    return matches.slice(0, limit).map((entry) => entry.word);
  }

  get wordCount() {
    return this.#wordCount;
  }

  #collect(node, prefix, matches) {
    if (node.terminal) {
      matches.push({ word: prefix, weight: node.weight });
    }
    for (const [character, child] of node.children) {
      this.#collect(child, prefix + character, matches);
    }
  }
}

const trie = new Trie();
for (const [word, weight] of [
  ["apple", 5],
  ["app", 9],
  ["apricot", 3],
  ["banana", 7],
  ["appetite", 4],
  ["application", 6],
]) {
  trie.insert(word, weight);
}
console.log(trie.suggest("app"));
console.log(trie.has("APP"), trie.startsWith("APR"), trie.wordCount);
console.log(trie.suggest("zz"));