haha-dict-tree/src/tree.js

因为最近常用到字典树，所以自己封装一个易用的字典树库。 代码量很少，90行，使用 generator 来遍历字典树，写起来很舒服。

class Node {
  constructor(char, word) {
    this.char = char;
    this.word = word;
    this.isWord = false;
    this.deep = word.length;
    this.number = 1;
    this.nextMap = new Map();
  }

  add(char, word) {
    const next = this.nextMap.get(char);
    if (next) {
      return next;
    }
    const n = new Node(char, word);
    this.nextMap.set(char, n);
    return n;
  }

  next(char) {
    return this.nextMap.get(char);
  }
}

class Tree {
  constructor() {
    this.root = new Node('root', '0');
  }

  addWord(word) {
    var current = this.root;
    var w = '';
    for (let c of word.split('')) {
      w = w + c;
      current = current.add(c, w);
    }
    current.isWord = true;
    current.number ++;
    return current;
  }

  *findAllRaw(word) {
    var current = this.root;
    for (let c of word.split('')) {
      current = current.next(c);
      if (!current) break;
      yield current;
    }

    if (current) {
      var nextMap = [current.nextMap];
      var temp = [];

      while (nextMap.length >= 1) {
        for (let m of nextMap) {
          for (let [k, v] of m.entries()) {
            void k;
            temp.push(v.nextMap);
            yield v;
          }
        }
        nextMap = temp;
        temp = [];
      }
    }
  }

  *findAll(word) {
    for (let n of this.findAllRaw(word)) {
      if (!n.isWord) continue;
      yield n;
    }
  }

  *findBefore(word) {
    for (const n of this.findAll(word)) {
      if (n.word.length > word.length) break;
      yield n;
    }
  }

  *findAfter(word) {
    for (const n of this.findAll(word)) {
      if (n.word.length < word.length) continue;
      yield n;
    }
  }
}

module.exports = Tree;