【刷穿 LeetCode】一题双解 :「朴素解法」&「二分 + 优先队列（堆）」

题目描述

这是 LeetCode 上的 1337. 矩阵中战斗力最弱的 K 行 ，难度为简单。

Tag : 「优先队列」、「二分」

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

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

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

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

示例 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

朴素解法

一个朴素的做法是对矩阵进行遍历，统计每一行的军人数量，并以二元组的形式进行存储。

然后对所有行的战力进行排序，选出战力最小的个下标即是答案。

代码：

class Solution {
    public int[] kWeakestRows(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        int[][] all = new int[m][1];
        for (int i = 0; i < m; i++) {
            int cur = 0;
            for (int j = 0; j < n; j++) cur += mat[i][j];
            all[i] = new int[]{cur, i};
        }
        Arrays.sort(all, (a, b)->{
            if (a[0] != b[0]) return a[0] - b[0];
            return a[1] - b[1];
        });
        int[] ans = new int[k];
        for (int i = 0; i < k; i++) ans[i] = all[i][1];
        return ans;
    }
}

时间复杂度：遍历矩阵的复杂度为；排序复杂度为；构造答案复杂度为。整体复杂度为
空间复杂度：空间用于存储所有的行战力；空间用于排序。整体复杂度为

二分 + 优先队列（堆）

根据「军人总是排在靠前位置」的特性，我们可以通过「二分」找到每一行最后一个军人的位置，从而在的复杂度内统计出每行的军人数量（战力情况）。

同时由于我们只需要「战力最弱」的行数据，这引导我们可以建立一个「战力大根堆」来做，「战力大根堆」存放着数量最多为的战力二元组，堆顶元素为战力最大的数对。

每次通过「二分」得到当前行的战力值后，判断当前战力值与堆顶元素的战力大小关系：

如果当前战力值比堆顶的元素要大：直接丢弃当前战力值（不可能属于在第小的集合中）；
如果当前战力值比堆顶的元素要小：将堆顶元素弹出，将当前行放入堆中。

一些细节：每次写二分都有同学会问，check 函数怎么写，可以重点看看 34 题题解。由于考虑某行没有军人的情况，我们需要二分完检查一下分割点是否符合 check 来决定军人数量。

另外，从堆弹出和往堆存入，需要与当前堆元素的数量有所关联。

代码：

class Solution {
    public int[] kWeakestRows(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        PriorityQueue<int[]> q = new PriorityQueue<>((a,b)->{
            if (a[0] != b[0]) return b[0] - a[0];
            return b[1] - a[1];
        });
        for (int i = 0; i < m; i++) {
            int l = 0, r = n - 1;
            while (l < r) {
                int mid = l + r + 1 >> 1;
                if (mat[i][mid] >= 1) l = mid;
                else r = mid - 1;
            }
            int cur = mat[i][r] >= 1 ? r + 1 : r;
            if (q.size() == k && q.peek()[0] > cur) q.poll();
            if (q.size() < k) q.add(new int[]{cur, i});
        }
        int[] ans = new int[k];
        int idx = k - 1;
        while (!q.isEmpty()) ans[idx--] = q.poll()[1];
        return ans;
    }
}