leetcode-go/medium/378.kth-smallest-element-in-a-sorted-matrix.go

type array [][2]int

func (a array) Len() int           { return len(a) }
func (a array) Less(i, j int) bool { return a[i][0] < a[j][0] }
func (a array) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func (a *array) Push(x interface{}) {
	*a = append(*a, x.([2]int))
}

func (a *array) Pop() interface{} {
	x := (*a)[a.Len()-1]
	*a = (*a)[:a.Len()-1]
	return x
}

func kthSmallest(matrix [][]int, k int) int {
	n := len(matrix)
	idx := make([]int, n)
	var pq array = make([][2]int, 0)
	for i := 0; i < n; i++ {
		heap.Push(&pq, [2]int{matrix[i][0], i})
	}
	for i := 0; i < k-1; i++ {
		x := heap.Pop(&pq).([2]int)
		idx[x[1]]++
		if nid := idx[x[1]]; nid < n {
			heap.Push(&pq, [2]int{matrix[x[1]][nid], x[1]})
		}
	}
	return heap.Pop(&pq).([2]int)[0]
}