sebinsua/interesting-trie.ts

export class TrieNode {
  public children: Record<string, TrieNode> = {};
  public toRelatedNodes: Set<TrieNode> = new Set();
  public isEndOfWord = false;
  public word?: string;
}

export class Trie {
  private root: TrieNode = new TrieNode();

  private getNode(prefix: string, isWord = false): TrieNode | undefined {
    let node = this.root;
    for (const char of prefix) {
      if (!node.children[char]) {
        return;
      }
      node = node.children[char];
    }
    return isWord ? (node.isEndOfWord ? node : undefined) : node;
  }

  add(word: string) {
    let node = this.root;
    for (const char of word) {
      if (!node.children[char]) {
        node.children[char] = new TrieNode();
      }
      node = node.children[char];
    }
    node.isEndOfWord = true;
    node.word = word;
  }

  has(word: string): boolean {
    let searchSpace = [this.root];
    for (const char of word) {
      searchSpace =
        char === "*"
          ? searchSpace.flatMap((node) => Object.values(node.children))
          : searchSpace.flatMap((node) =>
              node.children[char] ? [node.children[char]] : []
            );

      if (searchSpace.length === 0) {
        return false;
      }
    }

    return searchSpace.some((node) => node.isEndOfWord);
  }

  addRelation(word1: string, word2: string) {
    const node1 = this.getNode(word1, true);
    const node2 = this.getNode(word2, true);
    if (node1 && node2) {
      node1.toRelatedNodes.add(node2);
      node2.toRelatedNodes.add(node1);
    }
  }

  *autocomplete(prefix: string): Generator<string> {
    let nodesToVisit = [{ node: this.root, word: "" }];

    for (const char of prefix) {
      const nextNodes: Array<{ node: TrieNode; word: string }> = [];
      for (const { node, word } of nodesToVisit) {
        if (char === "*") {
          // Add all child nodes when encountering a wildcard
          for (const [childChar, childNode] of Object.entries(node.children)) {
            nextNodes.push({ node: childNode, word: word + childChar });
          }
        } else if (node.children[char]) {
          nextNodes.push({ node: node.children[char], word: word + char });
        }
      }
      nodesToVisit = nextNodes;
    }

    const seenWords = new Set<string>();
    const seenNodes = new Set<TrieNode>();
    while (nodesToVisit.length) {
      const { node, word } = nodesToVisit.shift()!;

      if (node.isEndOfWord && !seenWords.has(word)) {
        yield word;
        seenWords.add(word);
      }

      if (!seenNodes.has(node)) {
        seenNodes.add(node);
        for (const [childChar, childNode] of Object.entries(node.children)) {
          nodesToVisit.push({ node: childNode, word: word + childChar });
        }

        for (const relatedNode of node.toRelatedNodes) {
          if (!seenNodes.has(relatedNode)) {
            nodesToVisit.push({ node: relatedNode, word: relatedNode.word });
          }
        }
      }
    }
  }

  delete(word: string) {
    const nodesToRemove: Array<{ parent: TrieNode; childKey: string }> = [];

    let node = this.root;
    for (const char of word) {
      if (!node.children[char]) {
        return;
      }
      nodesToRemove.push({ parent: node, childKey: char });
      node = node.children[char];
    }
    if (!node.isEndOfWord) {
      return;
    }

    // Clear bi-directional relations.
    for (const relatedNode of node.toRelatedNodes) {
      relatedNode.toRelatedNodes.delete(node);
    }
    node.toRelatedNodes.clear();

    // Mark the node has no longer a word.
    node.isEndOfWord = false;

    // If there are nodes that don't form other words prune them.
    for (let i = nodesToRemove.length - 1; i >= 0; i--) {
      const { parent, childKey } = nodesToRemove[i];

      const childrenCount = Object.keys(
        parent.children[childKey].children
      ).length;
      if (parent.children[childKey].isEndOfWord || childrenCount > 0) {
        break;
      }

      delete parent.children[childKey];
    }
  }
}

const trie = new Trie();

trie.add("hello");
trie.add("hell");
trie.add("cello");
trie.add("hero");
trie.add("heroic");

trie.addRelation("hell", "hero");

console.log([...trie.autocomplete("hel")]);

trie.delete("hello");
console.log([...trie.autocomplete("*el")]);