Dustin Boston

/*!
 * @file Implements the Ternary Search Tree algorithm in TypeScript.
 * @author Dustin Boston <mail@dustin.boston>
 * @see Sedgewick, R., & Wayne, K. (2011). _Algorithms_ (4th ed.). Addison-Wesley.
 * @see Sedgewick, R., & Wayne, K. (n.d.). _TST.java._ Algorithms, 4th Ed.
 * Retrieved May 1, 2024, from https://algs4.cs.princeton.edu/52trie/TST.java.html
 */

/** Node class for Ternary Search Trie. */
export class TernarySearchTrieNode<Value> {
  character?: string;
  left?: TernarySearchTrieNode<Value>;
  mid?: TernarySearchTrieNode<Value>;
  right?: TernarySearchTrieNode<Value>;
  value?: Value;
}

/** Symbol table with string keys, implemented using a ternary search trie (TST). */
export class TernarySearchTrie<Value> {
  private nodeCount = 0;
  private rootNode?: TernarySearchTrieNode<Value>;

  /**
   * Returns the number of keys in the trie.
   *
   * @returns The size of the trie.
   */
  public size() {
    return this.nodeCount;
  }

  /**
   * Determines if a key is contained in the trie.
   *
   * @param key - The key to check.
   * @returns True if the key is contained, false otherwise.
   */
  public contains(key: string): boolean {
    if (!key) throw new Error("A key is required");
    return this.getValue(key) !== undefined;
  }

  /**
   * Retrieves the value associated with a given key.
   *
   * @param key - The key for which to retrieve the value.
   * @returns The value associated with the key, or undefined if not found.
   */
  public getValue(key: string): Value | undefined {
    if (!key || typeof key !== "string") {
      throw new Error("A string key is required");
    }

    const node = this.getNode(this.rootNode, key, 0);
    if (node === undefined) return;
    return node.value;
  }

  /**
   * Inserts or updates a key with a given value in the trie.
   *
   * @param key - The key to insert or update.
   * @param value - The value to associate with the key.
   */
  public put(key: string, value: Value): void {
    if (key === undefined) throw new Error("A key is required");

    if (!this.contains(key)) {
      this.nodeCount += 1;
    } else if (value === undefined) {
      this.nodeCount -= 1;
    }

    this.rootNode = this.putNode(this.rootNode, key, value, 0);
  }

  /**
   * Finds the longest prefix of the query string that is a key in the trie.
   *
   * @param query - The query string to search for the longest prefix.
   * @returns The longest prefix found, or undefined if none exists.
   */
  public longestPrefixOf(query: string): string | undefined {
    if (query === undefined) throw new Error("A query is required");
    if (query.length === 0) return;

    let length = 0;
    let node = this.rootNode;

    let index = 0;
    while (node?.character !== undefined && index < query.length) {
      const charCode = query.codePointAt(index) ?? 0;
      const nodeCharCode = node.character.codePointAt(0) ?? 0;

      if (charCode < nodeCharCode) {
        node = node.left;
      } else if (charCode > nodeCharCode) {
        node = node.right;
      } else {
        index++;
        if (node.value !== undefined) length = index;
        node = node.mid;
      }
    }

    return query.slice(0, Math.max(0, length));
  }

  /**
   * Returns all keys in the trie that match the pattern. The period character
   * (.) is a wildcard that represents any character. Doesn't match anything
   * longer than the given pattern.
   *
   * @param pattern - The pattern to match keys against, "." means any char.
   * @returns An array of keys that match the pattern.
   */
  public keysThatMatch(pattern: string): string[] {
    const keys: string[] = [];
    this.collectKeysThatMatch(this.rootNode, "", 0, pattern, keys);
    return keys;
  }

  /**
   * Recursive method to retrieve a node corresponding to the key.
   *
   * @param node - The current node in the trie.
   * @param key - The key to search for.
   * @param index - The index of the character in the key.
   * @returns The node corresponding to the key or undefined if not found.
   */
  private getNode(
    node: TernarySearchTrieNode<Value> | undefined,
    key: string | undefined,
    index: number,
  ): TernarySearchTrieNode<Value> | undefined {
    if (node?.character === undefined) return;
    if (!key) throw new Error("A key is required");

    const charCode = key.codePointAt(index) ?? 0;
    const nodeCharCode = node.character.codePointAt(0) ?? 0;
    if (charCode < nodeCharCode) {
      return this.getNode(node.left, key, index);
    }

    if (charCode > nodeCharCode) {
      return this.getNode(node.right, key, index);
    }

    if (index < key.length - 1) {
      return this.getNode(node.mid, key, index + 1);
    }

    return node;
  }

  /**
   * Recursive method to insert or update a node in the trie.
   *
   * @param node - The current node.
   * @param key - The key to insert or update.
   * @param value - The value to associate with the key.
   * @param index - The index of the character in the key.
   * @returns The updated or newly inserted node.
   */
  private putNode(
    node: TernarySearchTrieNode<Value> | undefined,
    key: string,
    value: Value,
    index: number,
  ): TernarySearchTrieNode<Value> {
    if (node?.character === undefined) {
      node = new TernarySearchTrieNode();
      node.character = key.charAt(index);
    }

    const charCode = key.codePointAt(index) ?? 0;
    const nodeCharCode = node.character.codePointAt(0) ?? 0;
    if (charCode < nodeCharCode) {
      node.left = this.putNode(node.left, key, value, index);
    } else if (charCode > nodeCharCode) {
      node.right = this.putNode(node.right, key, value, index);
    } else if (index < key.length - 1) {
      node.mid = this.putNode(node.mid, key, value, index + 1);
    } else {
      node.value = value;
    }

    return node;
  }

  /**
   * Collects keys recursively for the keys method.
   *
   * @param node - The current node.
   * @param prefix - The current prefix formed by traversing the trie.
   * @param keys - The array to collect keys into.
   */
  private collectKeys(
    node: TernarySearchTrieNode<Value> | undefined,
    prefix: string,
    keys: string[],
  ) {
    if (node === undefined) return;
    this.collectKeys(node.left, prefix, keys);
    if (node.value !== undefined) keys.push(prefix + node.character);
    prefix += node.character;
    this.collectKeys(node.mid, prefix, keys);
    prefix = prefix.slice(0, Math.max(0, prefix.length - 1));
    this.collectKeys(node.right, prefix, keys);
  }

  /**
   * Recursively collects keys that match the given pattern. The period
   * character (.) is a wildcard that represents any character. Doesn't match
   * anything longer than the given pattern.
   *
   * @param node - The current node.
   * @param prefix - The current prefix formed by traversing the trie.
   * @param index - The current index in the pattern.
   * @param pattern - The pattern to match keys against.
   * @param result - The array to collect matching keys into.
   */
  // eslint-disable-next-line max-params
  private collectKeysThatMatch(
    node: TernarySearchTrieNode<Value> | undefined,
    prefix: string,
    index: number,
    pattern: string,
    keys: string[],
  ) {
    if (node?.character === undefined) return;

    const patternLength = pattern.length - 1;
    const char = pattern.charAt(index);
    const charCode = char.codePointAt(0) ?? 0;
    const nodeChar = node.character;
    const nodeCharCode = nodeChar.codePointAt(0) ?? 0;

    if (char === "." || charCode < nodeCharCode) {
      this.collectKeysThatMatch(node.left, prefix, index, pattern, keys);
    }

    if (char === "." || charCode === nodeCharCode) {
      if (index === patternLength && node.value !== undefined) {
        keys.push(prefix + nodeChar);
      }

      if (index < patternLength) {
        prefix += nodeChar;
        this.collectKeysThatMatch(node.mid, prefix, index + 1, pattern, keys);
        prefix = prefix.slice(0, Math.max(0, prefix.length - 1));
      }
    }

    if (char === "." || charCode > nodeCharCode) {
      this.collectKeysThatMatch(node.right, prefix, index, pattern, keys);
    }
  }
}

// Usage
const wordList = "she sells sea shells by the sea shore".split(" ");

const tst = new TernarySearchTrie<number>();
const uniqueKeys = new Set<string>();
for (const [value, key] of wordList.entries()) {
  uniqueKeys.add(key);
  tst.put(key, value);
  console.log(`Put a new key: ${key}.`);
}

console.log(`Size: ${tst.size()}`);

const longestPrefix = tst.longestPrefixOf("shellsort");
console.log(`Longest prefix: ${longestPrefix}`);

import {TernarySearchTrie} from "../../../../data/algorithms/src/ternary_search_trie.ts";
import {assert} from "$assert";

const WORD_LIST = "she sells sea shells by the sea shore".split(" ");

function getTST() {
  const tst = new TernarySearchTrie<number>();
  const uniqueKeys = new Set<string>();
  for (const [value, key] of WORD_LIST.entries()) {
    uniqueKeys.add(key);
    tst.put(key, value);
  }

  return {tst, uniqueKeys};
}

test(function testPut() {
  const {tst, uniqueKeys} = getTST();
  const actual = tst.size();
  const expected = uniqueKeys.size;
  assert(actual === expected);
});

test(function testLongestPrefixOf() {
  const {tst} = getTST();
  const actual = tst.longestPrefixOf("shellsort");
  assert(actual === "shells");
});

test(function testKeys() {
  const {tst} = getTST();
  const actual = tst.keys().join(",");
  const expected = "by,sea,sells,she,shells,shore,the";
  assert(actual === expected);
});

test(function testKeysWithPrefix() {
  const {tst} = getTST();
  const actual = tst.keysWithPrefix("shor").join(",");
  assert(actual === "shore");
});

test(function testKeysThatMatch() {
  const {tst} = getTST();
  const actual = tst.keysThatMatch(".he.l.").join(",");
  assert(actual === "shells");
});