Skip to content

前缀树

TrieNode

  • trie/TrieNode.js
import HashTable from '../hash-table/HashTable.js';

export default class TrieNode {
  constructor(character, isCompleteWord = false) {
    this.character = character;
    this.isCompleteWord = isCompleteWord;
    this.children = new HashTable();
  }

  getChild(character) {
    return this.children.get(character);
  }

  hasChildren() {
    return this.children.getKeys().length !== 0;
  }

  addChild(character, isCompleteWord = false) {
    if (!this.children.has(character)) {
      this.children.set(character, new TrieNode(character, isCompleteWord));
    }

    const childNode = this.getChild(character);
    childNode.isCompleteWord = childNode.isCompleteWord || isCompleteWord;

    return childNode;
  }

  /**
   * 
   * @param {string} character 
   * @returns {TrieNode}
   */
  removeChild(character) {
    const childNode = this.getChild(character);

    // Delete childNode only if:
    // - childNode has NO children,
    // - childNode.isCompleteWord === false.
    if (
      childNode
      && !childNode.hasChildren()
      && !childNode.isCompleteWord
    ) {
      this.children.delete(character);
    }

    return this;
  }

  suggestChildren() {
    return [...this.children.getKeys()];
  }

  toString() {
    let childrenAsString = this.suggestChildren().toString();
    childrenAsString = childrenAsString ? `:${childrenAsString}` : ''
    const isCompleteString = this.isCompleteWord ? '*' : '';

    return `${this.character}${isCompleteString}${childrenAsString}`;
  }
}

Trie

  • trie/Trie.js
import TrieNode from './TrieNode.js';

const DEFAULT_CHARACTER = '*';

export default class Trie {
  constructor() {
    this.head = new TrieNode(DEFAULT_CHARACTER);
  }

  /**
   * 
   * @param {String} word 
   * @returns {Trie} 
   */
  addWord(word) {
    const characters = Array.from(word);
    let currentNode = this.head;

    for (let charIndex = 0; charIndex < characters.length; charIndex += 1) {
      const isComplete = charIndex === characters.length - 1;
      currentNode = currentNode.addChild(characters[charIndex], isComplete);
    }

    return this;
  }

  /**
   * Case:
   * trie.addWord('app').addWord('apple');
   * trie.deleteWord('app');
   * trie.doesWordExist('apple'); // true
   * 
   * @param {string} word
   * @return {Trie}
   */
  deleteWord(word) {
    const depthFirstDelete = (currentNode, charIndex = 0) => {
      if (charIndex >= word.length) {
        return;
      }

      const character = word[charIndex];
      const nextNode = currentNode.getChild(character);

      if (nextNode === null) {
        return;
      }

      depthFirstDelete(nextNode, charIndex + 1);

      // Since we're going to delete a word
      // let's un-mark its last character isCompleteWord flag.
      if (charIndex === (word.length - 1)) {
        nextNode.isCompleteWord = false;
      }

      // childNode is deleted only if:
      // - childNode has NO children
      // - childNode.isCompleteWord === false
      currentNode.removeChild(character);
    };

    depthFirstDelete(this.head);

    return this;
  }

  doesWordExist(word) {
    const lastCharacter = this.getLastCharacterNode(word);

    return !!lastCharacter && lastCharacter.isCompleteWord;
  }

  suggestNextCharacters(word) {
    const lastCharacter = this.getLastCharacterNode(word);

    if (!lastCharacter) {
      return null;
    }

    return lastCharacter.suggestChildren();
  }

  /**
   * @param {string} word
   * @return {TrieNode}
   */
  getLastCharacterNode(word) {
    const characters = Array.from(word);
    let currentNode = this.head;

    for (let charIndex = 0; charIndex < characters.length; charIndex += 1) {
      currentNode = currentNode.getChild(characters[charIndex]);
      if (!currentNode) {
        return null;
      }
    }

    return currentNode;
  }
}