Combining heap with hash table to get a data structure that supports fast iteration even with frequent insert and delete

OrderedCollections.ts

export class OrderedSet<T> {
  private heap: T[] = [];
  private sorted: T[] = [];
  private heapIndex = new Map<T, number>();
  private comparator: (item: T) => any;

  constructor(comparator: (item: T) => any = item => item) {
    this.comparator = comparator;
  }

  clear(): void {
    this.heap = [];
    this.sorted = [];
    this.heapIndex.clear();
  }

  add(item: T): boolean {
    if (this.sorted.length > 0) this.restoreHeap();
    if (this.heapIndex.has(item)) {
      const index = this.heapIndex.get(item)!;
      this.bubbleDown(index);
      this.bubbleUp(index);
      return false;
    }
    this.heap.push(item);
    const index = this.heap.length - 1;
    this.heapIndex.set(item, index);
    this.bubbleUp(index);
    return true;
  }

  pop(): T | undefined {
    if (this.heapIndex.size === 0) return undefined;
    if (this.sorted.length > 0) this.restoreHeap();
    return this.dequeue(0);
  }

  peek(): T | undefined {
    if (this.sorted.length > 0) return this.sorted[0];
    return this.heap[0];
  }

  delete(item: T): boolean {
    if (!this.heapIndex.has(item)) return false;
    if (this.sorted.length > 0) this.restoreHeap();
    const index = this.heapIndex.get(item)!;
    this.dequeue(index);
    return true;
  }

  has(item: T): boolean {
    return this.heapIndex.has(item);
  }

  at(index: number): T | undefined {
    let lastSortedIndex = this.sorted.length - 1;
    if (index <= lastSortedIndex) return this.sorted[index];
    while (this.heap.length > 0) {
      const min = this.dequeue(0);
      this.sorted.push(min);
      lastSortedIndex += 1;
      this.heapIndex.set(min, lastSortedIndex);
      if (index <= lastSortedIndex) return this.sorted[index];
    }
    return undefined;
  }

  private restoreHeap(): void {
    const other =
      this.heap.length <= this.sorted.length ? this.heap : this.sorted;
    if (other === this.heap) {
      this.heap = this.sorted;
    }
    for (const item of other) {
      this.heap.push(item);
      const index = this.heap.length - 1;
      this.heapIndex.set(item, index);
      this.bubbleUp(index);
    }
    this.sorted = [];
  }

  private dequeue(index: number): T {
    const item = this.heap[index];
    this.heapIndex.delete(item);
    const last = this.heap.pop()!;
    if (index < this.heap.length) {
      this.heap[index] = last;
      this.heapIndex.set(last, index);
      this.bubbleDown(index);
      this.bubbleUp(index);
    }
    return item;
  }

  private bubbleUp(cIndex: number): void {
    let curr = cIndex;
    while (curr > 0) {
      // eslint-disable-next-line no-bitwise
      const pIndex = (curr - 1) >> 1;
      if (this.swap(pIndex, curr)) {
        curr = pIndex;
      } else {
        break;
      }
    }
  }

  private bubbleDown(pIndex: number): void {
    const n = this.heap.length;
    let curr = pIndex;
    let cIndex = curr * 2 + 1;
    while (cIndex < n) {
      const altCIndex = cIndex + 1;
      if (altCIndex < n) {
        const child = this.heap[cIndex];
        const altChild = this.heap[altCIndex];
        if (this.comparator(altChild) < this.comparator(child)) {
          cIndex = altCIndex;
        }
      }
      if (this.swap(curr, cIndex)) {
        curr = cIndex;
        cIndex = curr * 2 + 1;
      } else {
        break;
      }
    }
  }

  private swap(pIndex: number, cIndex: number): boolean {
    const parent = this.heap[pIndex];
    const child = this.heap[cIndex];
    if (this.comparator(child) < this.comparator(parent)) {
      this.heap[pIndex] = child;
      this.heap[cIndex] = parent;
      this.heapIndex.set(child, pIndex);
      this.heapIndex.set(parent, cIndex);
      return true;
    }
    return false;
  }

  get size(): number {
    return this.heapIndex.size;
  }

  *[Symbol.iterator](): Iterator<T> {
    const n = this.size;
    for (let index = 0; index < n; index += 1) yield this.at(index)!;
  }

  // more performant for already sorted collections
  static fromOrdered<T>(
    iter: Iterable<T>,
    comparator: (item: T) => any = item => item
  ): OrderedSet<T> {
    const set = new OrderedSet(comparator);
    for (const item of iter) {
      set.sorted.push(item);
      set.heapIndex.set(item, set.sorted.length - 1);
    }
    return set;
  }
}

export class OrderedMap<K, V> {
  private map = new Map<K, V>();
  private internal: OrderedSet<K>;

  constructor(comparator: (value: V, key: K) => any = (_, key) => key) {
    this.internal = new OrderedSet(key => comparator(this.map.get(key)!, key));
  }

  clear(): void {
    this.map.clear();
    this.internal.clear();
  }

  get(key: K): V | undefined {
    return this.map.get(key);
  }

  set(key: K, value: V): boolean {
    this.map.set(key, value);
    return this.internal.add(key);
  }

  delete(key: K): boolean {
    this.map.delete(key);
    return this.internal.delete(key);
  }

  pop(): [K, V] | undefined {
    const key = this.internal.pop();
    return key && [key, this.map.get(key)!];
  }

  peek(): [K, V] | undefined {
    const key = this.internal.peek();
    return key && [key, this.map.get(key)!];
  }

  has(key: K): boolean {
    return this.internal.has(key);
  }

  at(index: number): [K, V] | undefined {
    const key = this.internal.at(index);
    return key && [key, this.map.get(key)!];
  }

  get size(): number {
    return this.internal.size;
  }

  *[Symbol.iterator](): Iterator<[K, V]> {
    for (const key of this.internal) {
      yield [key, this.map.get(key)!];
    }
  }

  // more performant for already sorted collections
  static fromOrdered<K, V>(
    iter: Iterable<[K, V]>,
    comparator: (value: V, key: K) => any
  ): OrderedMap<K, V> {
    const ordered = new OrderedMap(comparator);
    // eslint-disable-next-line
    // @ts-ignore
    const { sorted, heapIndex } = ordered.internal;
    for (const [key, value] of iter) {
      ordered.map.set(key, value);
      sorted.push(key);
      heapIndex.set(key, sorted.length - 1);
    }
    return ordered;
  }
}