题解 | 打印N个数组整体最大的Top K

package main

import (
	"bufio"
	"container/heap"
	"fmt"
	"os"
	"strconv"
)

type Elem struct {
	val int
	row int
	col int
}
type myheap []Elem

func (m myheap) Len() int           { return len(m) }
func (m myheap) Less(i, j int) bool { return m[i].val > m[j].val }
func (m myheap) Swap(i, j int)      { m[i], m[j] = m[j], m[i] }

func (m *myheap) Pop() interface{} {
	x := (*m)[len(*m)-1]
	(*m) = (*m)[:len(*m)-1]
	return x
}

func (m *myheap) Push(x interface{}) {
	*m = append(*m, x.(Elem))
}

func main() {
	for {
		scanner := bufio.NewScanner(os.Stdin)

		// 设置大缓冲区，防止大输入时出错
		buf := make([]byte, 1024*1024) // 1MB
		scanner.Buffer(buf, 1024*1024)
		scanner.Split(bufio.ScanWords) // 按单词扫描（空格、换行等分隔）

		// 读取第一行：n
		scanner.Scan()
		n, _ := strconv.Atoi(scanner.Text())
		if n == 0 {
			break
		}
		scanner.Scan()
		k, _ := strconv.Atoi(scanner.Text())

		arr := make([][]int, n)
		for i := 0; i < n; i++ {
			scanner.Scan()
			m, _ := strconv.Atoi(scanner.Text())
			arr[i] = make([]int, m)
			for j := 0; j < m; j++ {
				scanner.Scan()
				arr[i][j], _ = strconv.Atoi(scanner.Text())
			}
		}
		result := process(arr, k)
        for i:=0; i<len(result); i++ {
            fmt.Printf("%d ", result[i])
        }
	}
}

func process(arr [][]int, k int) []int {
	res := []int{}

	var my myheap
	heap.Init(&my)
	for i, r := range arr {
		if len(r) > 0 {
			heap.Push(&my, Elem{
				val: r[len(r)-1],
				row: i,
				col: len(r) - 1,
			})
		}
	}

   // 大根堆思路 完事
	for i := 0; i < k && my.Len() > 0; i++ {
		e := heap.Pop(&my).(Elem)
		if e.col-1 >= 0 {
			heap.Push(&my, Elem{
				val: arr[e.row][e.col-1],
				row: e.row,
				col: e.col - 1,
			})
		}
		res = append(res, e.val)
	}
	return res

}