Leetcode - The K Weakest Rows in a Matrix

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

Given a m * n matrix mat of ones (representing soldiers) and zeros (representing civilians), return the indexes of the k weakest rows in the matrix ordered from the weakest to the strongest.

A row i is weaker than row j, if the number of soldiers in row i is less than the number of soldiers in row j, or they have the same number of soldiers but i is less than j. Soldiers are always stand in the frontier of a row, that is, always ones may appear first and then zeros.

Example 1:

Input: mat = 
[[1,1,0,0,0],
 [1,1,1,1,0],
 [1,0,0,0,0],
 [1,1,0,0,0],
 [1,1,1,1,1]], 
k = 3
Output: [2,0,3]
Explanation: 
The number of soldiers for each row is: 
row 0 -> 2 
row 1 -> 4 
row 2 -> 1 
row 3 -> 2 
row 4 -> 5 
Rows ordered from the weakest to the strongest are [2,0,3,1,4]

I'm considering to use heap here. Build a heap and then pop k times.

type row struct {
	count int
	index int
}
type rowHeap []row

func (h rowHeap) Len() int { return len(h) }
func (h rowHeap) Less(i, j int) bool {
	if h[i].count < h[j].count {
		return true
	} else if h[i].count == h[j].count {
		if h[i].index < h[j].index {
			return true
		} else {
			return false
		}
	}
	return false
}
func (h rowHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

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

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

func kWeakestRows(mat [][]int, k int) []int {
	h := make(rowHeap, len(mat))
	for i, m := range mat {
		h[i] = row{
			count: sum(m),
			index: i,
		}
	}

	heap.Init(&h)
	result := []int{}
	for i := 0; i < k; i++ {
		result = append(result, heap.Pop(&h).(row).index)
	}

	return result
}

func sum(num []int) int {
	result := 0
	for _, n := range num {
		result += n
	}
	return result
}

• 相关文章
• 全局相关文章