bitburner-src/src/utils/Heap.ts

134 lines
3.9 KiB
TypeScript
Raw Normal View History

2022-04-01 23:52:47 +02:00
/** Binary heap. */
abstract class BinHeap<T> {
/**
* Heap data array consisting of [weight, payload] pairs, arranged by weight
* to satisfy heap condition.
*
* Encodes the binary tree by storing tree root at index 0 and
* left child of element i at `i * 2 + 1` and
* right child of element i at `i * 2 + 2`.
*/
protected data: [number, T][];
constructor() {
this.data = [];
}
/** Get number of elements in the heap. */
2022-04-02 04:16:48 +02:00
public get size(): number {
2022-04-01 23:52:47 +02:00
return this.data.length;
}
2022-04-02 04:16:48 +02:00
/** Add a new element to the heap. */
public push(value: T, weight: number): void {
const i = this.data.length;
this.data[i] = [weight, value];
this.heapifyUp(i);
}
/** Get the value of the root-most element of the heap, without changing the heap. */
public peek(): T | undefined {
2022-04-08 06:39:45 +02:00
if (this.data.length == 0) return undefined;
2022-04-02 04:16:48 +02:00
return this.data[0][1];
}
/** Remove the root-most element of the heap and return the removed element's value. */
public pop(): T | undefined {
2022-04-08 06:39:45 +02:00
if (this.data.length == 0) return undefined;
2022-04-02 04:16:48 +02:00
const value = this.data[0][1];
this.data[0] = this.data[this.data.length - 1];
this.data.length = this.data.length - 1;
this.heapifyDown(0);
return value;
}
/** Change the weight of an element in the heap. */
public changeWeight(predicate: (value: T) => boolean, weight: number): void {
// Find first element with matching value, if any
2022-04-08 06:39:45 +02:00
const i = this.data.findIndex((e) => predicate(e[1]));
if (i == -1) return;
2022-04-02 04:16:48 +02:00
// Update that element's weight
this.data[i][0] = weight;
// And re-heapify if needed
const p = Math.floor((i - 1) / 2);
2022-04-08 06:39:45 +02:00
if (!this.heapOrderABeforeB(this.data[p][0], this.data[i][0]))
// Needs to shift root-wards?
2022-04-02 04:16:48 +02:00
this.heapifyUp(i);
2022-04-08 06:39:45 +02:00
// Try shifting deeper
else this.heapifyDown(i);
2022-04-02 04:16:48 +02:00
}
2022-04-01 23:52:47 +02:00
/** Restore heap condition, starting at index i and traveling towards root. */
protected heapifyUp(i: number): void {
// Swap the new element up towards root until it reaches root position or
// settles under under a suitable parent
2022-04-08 06:39:45 +02:00
while (i > 0) {
2022-04-01 23:52:47 +02:00
const p = Math.floor((i - 1) / 2);
// Reached heap-ordered state already?
2022-04-08 06:39:45 +02:00
if (this.heapOrderABeforeB(this.data[p][0], this.data[i][0])) break;
2022-04-01 23:52:47 +02:00
// Swap
const tmp = this.data[p];
this.data[p] = this.data[i];
this.data[i] = tmp;
// And repeat at parent index
i = p;
}
}
/** Restore heap condition, starting at index i and traveling away from root. */
protected heapifyDown(i: number): void {
// Swap the shifted element down in the heap until it either reaches the
// bottom layer or is in correct order relative to it's children
2022-04-08 06:39:45 +02:00
while (i < this.data.length) {
2022-04-01 23:52:47 +02:00
const l = i * 2 + 1;
const r = i * 2 + 2;
let toSwap = i;
// Find which one of element i and it's children should be closest to root
2022-04-08 06:39:45 +02:00
if (l < this.data.length && this.heapOrderABeforeB(this.data[l][0], this.data[toSwap][0])) toSwap = l;
if (r < this.data.length && this.heapOrderABeforeB(this.data[r][0], this.data[toSwap][0])) toSwap = r;
2022-04-01 23:52:47 +02:00
// Already in order?
2022-04-08 06:39:45 +02:00
if (i == toSwap) break;
2022-04-01 23:52:47 +02:00
// Not in order. Swap child that should be closest to root up to 'i' and repeat
const tmp = this.data[toSwap];
this.data[toSwap] = this.data[i];
this.data[i] = tmp;
i = toSwap;
}
}
2022-04-02 04:16:48 +02:00
/**
* Should element with weight `weightA` be closer to root than element with
* weight `weightB`?
*/
protected abstract heapOrderABeforeB(weightA: number, weightB: number): boolean;
2022-04-01 23:52:47 +02:00
}
/** Binary max-heap. */
export class MaxHeap<T> extends BinHeap<T> {
heapOrderABeforeB(weightA: number, weightB: number): boolean {
return weightA > weightB;
}
}
/** Binary min-heap. */
export class MinHeap<T> extends BinHeap<T> {
heapOrderABeforeB(weightA: number, weightB: number): boolean {
return weightA < weightB;
}
}