@thi.ng/trie/trie-map.js

Trie-based ES6-like Map data structures with prefix search/query support

class TrieMap {
  next = {};
  val;
  n = 0;
  constructor(pairs) {
    pairs && this.into(pairs);
  }
  get size() {
    let size = 0;
    const stack = [this];
    while (stack.length) {
      const node = stack.pop();
      if (node.val !== void 0) size++;
      if (node.next) stack.push(...Object.values(node.next));
    }
    return size;
  }
  [Symbol.iterator]() {
    return this.iterate((key, node) => [key, node.val]);
  }
  keys(prefix, includePrefix = true) {
    return this.iterate((key) => key, prefix, includePrefix);
  }
  values(prefix) {
    return this.iterate((_, node) => node.val, prefix);
  }
  entries(prefix, includePrefix = true) {
    return this.iterate(
      (key, node) => [key, node.val],
      prefix,
      includePrefix
    );
  }
  clear() {
    this.next = {};
    this.n = 0;
    this.val = void 0;
  }
  has(key) {
    return this.get(key) !== void 0;
  }
  hasPrefix(prefix) {
    return !!this.find(prefix);
  }
  get(key, notFound) {
    const node = this.find(key);
    return node ? node.val : notFound;
  }
  find(key) {
    let node = this;
    for (let i = 0, n = key.length; i < n; i++) {
      node = node.next[key[i]];
      if (!node) return;
    }
    return node;
  }
  /**
   * Returns longest known prefix for `key`. Returns undefined if given key
   * has no partial matches.
   *
   * @param key -
   */
  knownPrefix(key) {
    let node = this;
    let prefix = "";
    for (let i = 0, n = key.length; i < n; i++) {
      const k = key[i];
      const next = node.next[k];
      if (!next) break;
      prefix += k;
      node = next;
    }
    return prefix || void 0;
  }
  hasKnownPrefix(key) {
    return !!this.knownPrefix(key);
  }
  set(key, val) {
    let node = this;
    for (let i = 0, n = key.length; i < n; i++) {
      const k = key[i];
      const next = node.next[k];
      node = !next ? (node.n++, node.next[k] = new TrieMap()) : next;
    }
    node.val = val;
  }
  into(pairs) {
    for (const [k, v] of pairs) {
      this.set(k, v);
    }
  }
  delete(prefix) {
    const n = prefix.length;
    if (n < 1) return false;
    const path = [];
    const key = [];
    let i = 0;
    let node = this;
    for (; i < n; i++) {
      const k = prefix[i];
      key.push(k);
      path.push(node);
      node = node.next[k];
      if (!node) return false;
    }
    while (node = path[--i]) {
      delete node.next[key[i]];
      if (--node.n) break;
    }
    return true;
  }
  *iterate(fn, prefix = "", includePrefix = true) {
    const root = this.find(prefix);
    if (!root) return;
    const stack = [
      [includePrefix ? prefix : "", root]
    ];
    while (stack.length) {
      const [key, node] = stack.pop();
      if (node.val !== void 0) yield fn(key, node);
      const next = node.next;
      for (let k in next) {
        stack.push([key + k, next[k]]);
      }
    }
  }
}
const defTrieMap = (pairs) => new TrieMap(pairs);
const defTrieMapFromJSON = (src) => {
  const res = defTrieMap();
  const queue = [["", src]];
  while (queue.length) {
    const [key, node] = queue.pop();
    if (node.val) res.set(key, node.val);
    for (const [k, child] of Object.entries(node.next)) {
      queue.push([key + k, child]);
    }
  }
  return res;
};
export {
  TrieMap,
  defTrieMap,
  defTrieMapFromJSON
};