Ternary Search Trie | The Personal Website of Dustin Boston

A Ternary Search Trie (TST) is a data structure used to store and retrieve strings efficiently, combining features of binary search trees and tries. It organizes data in a tree where each node contains three children (low, equal, high) corresponding to characters less than, equal to, or greater than the node’s character, enabling balanced search operations while conserving memory compared to standard tries.

code.ts

/*!
 * @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}`);

ternary_search_trie_test.ts

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");
});