Merge K sorted arrays in a single one. Link: https://goplay.space/#rDyKPxaNLcl

merge_k_sorted_arrays.go

package main

import (
	"container/heap"
	"fmt"
)

type Count struct {
	arrayRef int
	indexRef int
	value    int
}

// An IntHeap is a min-heap of ints.
type CountHeap []Count

func (h CountHeap) Len() int           { return len(h) }
func (h CountHeap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h CountHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *CountHeap) Push(x interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(Count))
}

func (h *CountHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// This example inserts several ints into an IntHeap, checks the minimum,
// and removes them in order of priority.
func main() {
	input := [][]int{
		{1, 5},
		{1, 3, 4, 7},
		{2, 6, 8},
	}
	result := []int{}

	// Initialize heap with first elements of each array
	h := &CountHeap{}
	for index, array := range input {
		*h = append(*h, Count{
			arrayRef: index,
			indexRef: 0,
			value:    array[0],
		})
	}
	heap.Init(h)

	for h.Len() > 0 {
		// Get minimum value
		min := heap.Pop(h).(Count)
		result = append(result, min.value)

		// Get array where we will extract next number
		selectedArray := input[min.arrayRef]

		// If next number exist, we push it to the heap
		if min.indexRef < len(selectedArray)-1 {
			heap.Push(h,
				Count{
					arrayRef: min.arrayRef,
					indexRef: min.indexRef + 1,
					value:    selectedArray[min.indexRef+1],
				},
			)
		}

	}

	fmt.Print(result)
}