Merge K sorted lists

merge.go

package main

import (
	"container/heap"
)

type priorityQueue [][]int

func (pq priorityQueue) Len() int { return len(pq) }

func (pq priorityQueue) Less(i, j int) bool {
	return pq[i][0] < pq[j][0]
}

func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *priorityQueue) Push(x any) {
	*pq = append(*pq, x.([]int))
}

func (pq *priorityQueue) Pop() any {
	old := *pq
	n := len(old)
	ret := old[n-1]
	old[n-1] = nil // avoid memory leak
	*pq = old[:n-1]
	return ret
}

func merge(lists ...[]int) []int {
	pq := make(priorityQueue, 0, len(lists))
	var outLen int

	for _, list := range lists {
		if len(list) == 0 {
			continue
		}
		pq = append(pq, list)
		outLen += len(list)
	}
	if outLen == 0 {
		return nil
	}

	heap.Init(&pq)
	merged := make([]int, 0, outLen)

	// Merge heads from lists until only one list in pq.
	for pq.Len() != 1 {
		x := heap.Pop(&pq)
		list := x.([]int)
		merged = append(merged, list[0])
		if len(list) == 1 {
			continue
		}
		heap.Push(&pq, list[1:])
	}

	// Merge remaining items from remaining list.
	x := heap.Pop(&pq)
	list := x.([]int)
	return append(merged, list...)
}