【二分查找】leetcode 1337. 矩阵中战斗力最弱的 K 行

1337. 矩阵中战斗力最弱的 K 行

题目描述

给你一个大小为 m * n 的矩阵 mat，矩阵由若干军人和平民组成，分别用 1 和 0 表示。

请你返回矩阵中战斗力最弱的 k 行的索引，按从最弱到最强排序。

如果第 i 行的军人数量少于第 j 行，或者两行军人数量相同但 i 小于 j，那么我们认为第 i 行的战斗力比第 j 行弱。

军人 总是 排在一行中的靠前位置，也就是说 1 总是出现在 0 之前。

示例1：

输入： 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
输出： [2,0,3]
解释： 每行中的军人数目：
行 0 -> 2
行 1 -> 4
行 2 -> 1
行 3 -> 2
行 4 -> 5
从最弱到最强对这些行排序后得到 [2,0,3,1,4]

示例2：

输入： mat =
[[1,0,0,0],
[1,1,1,1],
[1,0,0,0],
[1,0,0,0]],
k = 2
输出： [0,2]
解释：
每行中的军人数目：
行 0 -> 1
行 1 -> 4
行 2 -> 1
行 3 -> 1
从最弱到最强对这些行排序后得到 [0,2,3,1]

提示

• m == mat.length
• n == mat[i].length
• 2 <= n, m <= 100
• 1 <= k <= m
• matrix[i][j] 不是 0 就是 1

方法一：朴素做法

解题思路

统计出每行军人的数目，然后对其进行排序，输出前 k 行的索引 。

代码

class Solution {
typedef pair<int, int> PII;
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        vector<int> ans(k);
        int n = mat.size();
        vector<PII> cnt(n);
        for(int i = 0; i < mat.size(); i++)
        {
            cnt[i] = {count(mat[i].begin(), mat[i].end(), 1), i};
        }
        sort(cnt.begin(), cnt.end(), [&](const auto& x, const auto& y){
            return x.first == y.first ? x.second < y.second : x.first < y.first;
        });
        for(int i = 0; i < k; i++)
            ans[i] = cnt[i].second;
        return ans;
    }
};

复杂度分析

• 时间复杂度：遍历矩阵的复杂度为 $O(m\ast n)$；排序复杂度为 $O(mlogm)$；构造答案复杂度为 $O(k)$。整体复杂度为 $O(\max (m\ast n,m\log m))$。
• 空间复杂度：$O(m)$ 空间用于存储所有的行战力；$O(logm)$ 空间用于排序。整体复杂度为 $O(m+logm)$。

方法二：二分查找 + 堆

解题思路

军人总是排在一行中的靠前位置，也就是说 1 总是出现在 0 之前。因此，我们用二分查找来找到每一行中 1 出现的最后一个位置，下标 + 1 便是每行中 1 的个数，即每行的战斗力。
当我们得到每一行的战斗力后，我们可以将它们全部放入一个小根堆中，并不断地取出堆顶的元素 k 次，这样我们就得到了最弱的 k 行的索引。
如果我们依次将每一行的战斗力以及索引（因为如果战斗力相同，索引较小的行更弱，所以我们需要在小根堆中存放战斗力和索引的二元组）放入小根堆中，那么这样做的时间复杂度是 $O(m\log m)$ 的。
一种更好的方法是使用这 m 个战斗力值直接初始化一个小根堆，时间复杂度为 O(m)。即 $priority\_queueq(greater<pair<int,int>>(),move(power))$

代码

class Solution {
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        int m = mat.size(), n = mat[0].size();
        vector<pair<int, int>> power;
        for(int i = 0; i < m; i++)
        {
            int l = 0, r = n - 1, mid;
            while(l < r)
            {
                mid = (l + r + 1) / 2;
                if(mat[i][mid] >= 1)    l = mid;
                else    r = mid - 1;
            }
            power.emplace_back(l + 1, i);
        }
        // 数组直接初始化小根堆
        priority_queue q(greater<pair<int, int>>(), move(power));
        vector<int> ans;
        for(int i = 0; i < k; i++)
        {
            ans.push_back(q.top().second);
            q.pop();
        }
        return ans;
    }
};

复杂度分析

• 时间复杂度：O($mlogn+klogm$)。
  需要 $O(m\log n)$ 的时间对每一行进行二分查找。
  需要 O(m) 的时间建立小根堆。
  需要 $O(k\log m)$ 的时间从堆中取出 k 个最小的元素。
• 空间复杂度：O(m)，即为堆所需要的空间。