1337. The K Weakest Rows in a Matrix*

1337. 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]

Example 2:

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

Constraints:

• m == mat.length
• n == mat[i].length
• 2 <= n, m <= 100
• 1 <= k <= m
• matrix[i][j] is either 0 or 1.

C++ 实现 1

明确几个点, 首先越 weak 的 row, 它的 1 的数目越少; 如果两行的 1 个数相同, 那么排在前面的越 weak. 如果我使用最大堆存储 pair<int, int>, 其中第一个值表示 1 的数目, 用来表示 weak 的程度, 第二个值表示 row 的索引, 使用最大堆, 保存 k 个最 weak 的结果.

另外, 需要注意的是, 最大堆

priority_queue<pair<int, int>, vector<pair<int, int>>, std::less<pair<int, int>>> max_heap;

它对于其中 pair 的排序结果, 应该是:

p.first < q.first || (p.first == q.first && p.second < q.second)

只有这样, 最后输出的结果才可能正确, 因为当两行的 weak 程度相同时, 后一行相对更 stonger.

class Solution {
private:
	// 使用二分法找每行的第一个 0
    int find_first_zero(const vector<int> &nums) {
        int l = 0, r = nums.size() - 1;
        while (l <= r) {
            int mid = l + (r - l) / 2;
            if (nums[mid] == 1) l = mid + 1;
            else r = mid - 1;
        }
        return l >= nums.size() ? -1 : l;
    }
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        priority_queue<pair<int, int>, vector<pair<int, int>>, std::less<pair<int, int>>> max_heap;
        int m = mat.size(), n = mat[0].size();
        for (int i = 0; i < m; ++ i) {
            auto idx = find_first_zero(mat[i]);
            int weight = idx == -1 ? n : idx;
            max_heap.push(make_pair(weight, i));
            if (max_heap.size() > k) max_heap.pop();
        }
        vector<int> res;
        while (!max_heap.empty()) {
            res.push_back(max_heap.top().second);
            max_heap.pop();
        }
        std::reverse(res.begin(), res.end());
        return res;
    }
};

C++ 实现 2

代码更为简洁, 来自 LeetCode Submission; 使用排序的思路, 也不考虑什么二分法.

class Solution {
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        vector<pair<int,int>> qnt ;
        for ( int i=0 ; i<mat.size() ; i++ ) {
            int c = 0 ;
            for ( int j=0 ; j<mat[i].size() ; j++ ) {
                if(mat[i][j]==1)  c++ ;
            }
            qnt.push_back(pair<int,int>(c,i)) ;
        }
        std::sort(qnt.begin(),qnt.end()) ;
        vector<int> ret ;
        for ( auto it = qnt.begin() ; it!=qnt.end() ; ++it ) {
            k-- ;
            ret.push_back(it->second) ;
            if(k==0)    break ;
        }
    	return ret ;
    }
};