tools/core/stdlib/src/structs/BinaryHeap/index.ts

import { first } from '../../arrays';
import { isArray } from '../../types';
import type { BinaryHeapLike, Comparator } from './types';

export type { BinaryHeapLike, Comparator } from './types';

export interface BinaryHeapOptions<T> {
  comparator?: Comparator<T>;
}

/**
 * Default min-heap comparator for numeric values
 *
 * @param {number} a First element
 * @param {number} b Second element
 * @returns {number} Negative if a < b, positive if a > b, zero if equal
 */
const defaultComparator: Comparator<any> = (a: number, b: number) => a - b;

/**
 * @name BinaryHeap
 * @category Data Structures
 * @description Binary heap backed by a flat array with configurable comparator
 *
 * @since 0.0.8
 *
 * @template T The type of elements stored in the heap
 */
export class BinaryHeap<T> implements BinaryHeapLike<T> {
  /**
     * The comparator function used to order elements
     *
     * @private
     * @type {Comparator<T>}
     */
  private readonly comparator: Comparator<T>;

  /**
     * Internal flat array backing the heap
     *
     * @private
     * @type {T[]}
     */
  private readonly heap: T[] = [];

  /**
     * Creates an instance of BinaryHeap
     *
     * @param {(T[] | T)} [initialValues] The initial values to heapify
     * @param {BinaryHeapOptions<T>} [options] Heap configuration
     */
  constructor(initialValues?: T[] | T, options?: BinaryHeapOptions<T>) {
    this.comparator = options?.comparator ?? defaultComparator;

    if (initialValues !== null && initialValues !== undefined) {
      const items = isArray(initialValues) ? initialValues : [initialValues];
      this.heap.push(...items);
      this.heapify();
    }
  }

  /**
     * Gets the number of elements in the heap
     * @returns {number} The number of elements in the heap
     */
  public get length(): number {
    return this.heap.length;
  }

  /**
     * Checks if the heap is empty
     * @returns {boolean} `true` if the heap is empty, `false` otherwise
     */
  public get isEmpty(): boolean {
    return this.heap.length === 0;
  }

  /**
     * Pushes an element into the heap
     * @param {T} element The element to insert
     */
  public push(element: T): void {
    this.heap.push(element);
    this.siftUp(this.heap.length - 1);
  }

  /**
     * Removes and returns the root element (min or max depending on comparator)
     * @returns {T | undefined} The root element, or `undefined` if the heap is empty
     */
  public pop(): T | undefined {
    if (this.heap.length === 0) return undefined;

    const root = first(this.heap)!;
    const last = this.heap.pop()!;

    if (this.heap.length > 0) {
      this.heap[0] = last;
      this.siftDown(0);
    }

    return root;
  }

  /**
     * Returns the root element without removing it
     * @returns {T | undefined} The root element, or `undefined` if the heap is empty
     */
  public peek(): T | undefined {
    return first(this.heap);
  }

  /**
     * Removes all elements from the heap
     * @returns {this} The heap instance for chaining
     */
  public clear(): this {
    this.heap.length = 0;
    return this;
  }

  /**
     * Returns a shallow copy of the heap elements as an array (heap order, not sorted)
     * @returns {T[]} Array of elements in heap order
     */
  public toArray(): T[] {
    return this.heap.slice();
  }

  /**
     * Returns a string representation of the heap
     * @returns {string} String representation
     */
  public toString(): string {
    return `BinaryHeap(${this.heap.length})`;
  }

  /**
     * Iterator over heap elements in heap order
     */
  public* [Symbol.iterator](): Iterator<T> {
    yield* this.heap;
  }

  /**
     * Async iterator over heap elements in heap order
     */
  public async* [Symbol.asyncIterator](): AsyncIterator<T> {
    for (const element of this.heap)
      yield element;
  }

  /**
     * Restores heap property by sifting an element up
     *
     * @private
     * @param {number} index The index of the element to sift up
     */
  private siftUp(index: number): void {
    const heap = this.heap;
    const cmp = this.comparator;

    while (index > 0) {
      const parent = (index - 1) >> 1;

      if (cmp(heap[index]!, heap[parent]!) >= 0) break;

      const temp = heap[index]!;
      heap[index] = heap[parent]!;
      heap[parent] = temp;

      index = parent;
    }
  }

  /**
     * Restores heap property by sifting an element down
     *
     * @private
     * @param {number} index The index of the element to sift down
     */
  private siftDown(index: number): void {
    const heap = this.heap;
    const cmp = this.comparator;
    const length = heap.length;

    while (true) {
      let smallest = index;
      const left = 2 * index + 1;
      const right = 2 * index + 2;

      if (left < length && cmp(heap[left]!, heap[smallest]!) < 0) {
        smallest = left;
      }

      if (right < length && cmp(heap[right]!, heap[smallest]!) < 0) {
        smallest = right;
      }

      if (smallest === index) break;

      const temp = heap[index]!;
      heap[index] = heap[smallest]!;
      heap[smallest] = temp;

      index = smallest;
    }
  }

  /**
     * Builds heap from unordered array in O(n) using Floyd's algorithm
     *
     * @private
     */
  private heapify(): void {
    for (let i = (this.heap.length >> 1) - 1; i >= 0; i--) {
      this.siftDown(i);
    }
  }
}