632. Smallest Range Covering Elements from K Lists

import "container/heap"

type listHead struct {
	value     int
	listIndex int
	nextIndex int
}

type minByValue []listHead

func (items minByValue) Len() int { return len(items) }
func (items minByValue) Less(indexA, indexB int) bool {
	return items[indexA].value < items[indexB].value
}
func (items minByValue) Swap(indexA, indexB int) { items[indexA], items[indexB] = items[indexB], items[indexA] }

func (items *minByValue) Push(value any) {
	*items = append(*items, value.(listHead))
}

func (items *minByValue) Pop() any {
	old := *items
	lastIndex := len(old) - 1
	popped := old[lastIndex]
	*items = old[:lastIndex]
	return popped
}

func smallestRange(lists [][]int) []int {
	active := &minByValue{}
	heap.Init(active)
	maxValue := 0
	for listIndex, numbers := range lists {
		firstValue := numbers[0]
		heap.Push(active, listHead{
			value:     firstValue,
			listIndex: listIndex,
			nextIndex: 1,
		})
		if listIndex == 0 || firstValue > maxValue {
			maxValue = firstValue
		}
	}
	// invariant: heap root is the current minimum among the k chosen front elements
	bestLow := (*active)[0].value
	bestHigh := maxValue
	for {
		smallest := heap.Pop(active).(listHead)
		if smallest.nextIndex >= len(lists[smallest.listIndex]) {
			break
		}
		nextValue := lists[smallest.listIndex][smallest.nextIndex]
		if nextValue > maxValue {
			maxValue = nextValue
		}
		heap.Push(active, listHead{
			value:     nextValue,
			listIndex: smallest.listIndex,
			nextIndex: smallest.nextIndex + 1,
		})
		newMin := (*active)[0].value
		if maxValue-newMin < bestHigh-bestLow {
			bestLow = newMin
			bestHigh = maxValue
		}
	}
	return []int{bestLow, bestHigh}
}

class MinHeap {
	constructor() {
		this.values = [];
	}
	push(entry) {
		this.values.push(entry);
		this.bubbleUp(this.values.length - 1);
	}
	pop() {
		const top = this.values[0];
		const last = this.values.pop();
		if (this.values.length > 0) {
			this.values[0] = last;
			this.bubbleDown(0);
		}
		return top;
	}
	peek() {
		return this.values[0];
	}
	get size() {
		return this.values.length;
	}
	bubbleUp(index) {
		while (index > 0) {
			const parentIndex = Math.floor((index - 1) / 2);
			if (this.values[parentIndex].value <= this.values[index].value) {
				break;
			}
			[this.values[parentIndex], this.values[index]] = [
				this.values[index],
				this.values[parentIndex],
			];
			index = parentIndex;
		}
	}
	bubbleDown(index) {
		const length = this.values.length;
		while (true) {
			const leftIndex = index * 2 + 1;
			const rightIndex = index * 2 + 2;
			let smallest = index;
			if (
				leftIndex < length &&
				this.values[leftIndex].value < this.values[smallest].value
			) {
				smallest = leftIndex;
			}
			if (
				rightIndex < length &&
				this.values[rightIndex].value < this.values[smallest].value
			) {
				smallest = rightIndex;
			}
			if (smallest === index) {
				break;
			}
			[this.values[index], this.values[smallest]] = [
				this.values[smallest],
				this.values[index],
			];
			index = smallest;
		}
	}
}

function smallestRange(lists) {
	const heap = new MinHeap();
	let maxValue = Number.NEGATIVE_INFINITY;
	for (let listIndex = 0; listIndex < lists.length; listIndex++) {
		const firstValue = lists[listIndex][0];
		heap.push({
			value: firstValue,
			listIndex,
			nextIndex: 1,
		});
		maxValue = Math.max(maxValue, firstValue);
	}
	let bestLow = heap.peek().value;
	let bestHigh = maxValue;
	while (true) {
		const smallest = heap.pop();
		if (smallest.nextIndex >= lists[smallest.listIndex].length) {
			break;
		}
		const nextValue = lists[smallest.listIndex][smallest.nextIndex];
		maxValue = Math.max(maxValue, nextValue);
		heap.push({
			value: nextValue,
			listIndex: smallest.listIndex,
			nextIndex: smallest.nextIndex + 1,
		});
		const newMin = heap.peek().value;
		if (maxValue - newMin < bestHigh - bestLow) {
			bestLow = newMin;
			bestHigh = maxValue;
		}
	}
	return [bestLow, bestHigh];
}