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] }