https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/

Leetcode: 1337. The K Weakest Rows in a Matrix

func kWeakestRows(mat [][]int, k int) []int {
	// Uses the known constraint that the number of rows is <= 100:
	// Row score is a combination of soldier count
	// and the row index, where soldier count is the major base and row index 
	// is the minor base. When sorted, rows with matching
	// soldier counts will be differentiated and sorted appropriately by their index, which can easily be
	// obtained later with mod 100. 
	scoreRow := func(index int, row []int) int {
		count := 0
		for _, i := range row {
			if i == 0 {
				break
			}
			count++
		}
		return count * 100 + index
	}

	h := NewIntHeap(len(mat), 0, false)
	for i, row := range mat {
		heap.Push(h, scoreRow(i, row))
	}

	// pop the top k, calculate index
	weakestRows := make([]int, k)
	for i := 0; i < k; i++ {
		score := heap.Pop(h).(int)
		index := score % 100
		weakestRows[i] = index
	}
	return weakestRows
}

// Helper IntHeap

func NewIntHeap(initialSize int, maxSize int, desc bool) *IntHeap {
	return &IntHeap{
		data:    make([]int, initialSize)[0:0],
		desc:    desc,
		maxSize: maxSize,
	}
}

type IntHeap struct {
	data []int
	desc bool // true = descending, false = ascending
	initialSize int
	maxSize int
}

func (h *IntHeap) resetData() {
	h.data = make([]int, h.initialSize)[0:0]
	heap.Init(h)
}

func (h IntHeap) Len() int {
	return len(h.data)
}

func (h IntHeap) Less(i, j int) bool {
	if h.desc {
		return h.data[i] > h.data[j]
	}
	return h.data[i] < h.data[j]
}

func (h IntHeap) Swap(i, j int) {
	h.data[i], h.data[j] = h.data[j], h.data[i]
}

func (h *IntHeap) Push(x interface{}) {
	h.data = append(h.data, x.(int))
}

func (h *IntHeap) Prune() {
	if h.maxSize == 0 || h.Len() < h.maxSize {
		return
	}
	// pop the desired values
	r := make([]int, h.maxSize)
	for i := 0; i < h.maxSize; i++ {
		r[i] = heap.Pop(h).(int)
	}
	// reset heap data
	h.resetData()
	// push back
	for _, v := range r {
		heap.Push(h, v)
	}
}

func (h *IntHeap) Pop() interface{} {
	n := h.Len()
	x := h.data[n-1]
	h.data = h.data[0:n-1]
	return x
}