题目描述
这是 LeetCode 上的 1337. 矩阵中战斗力最弱的 K 行 ,难度为 简单。
Tag : 「优先队列」、「二分」
给你一个大小为 m * nm∗n 的矩阵 matmat,矩阵由若干军人和平民组成,分别用 11 和 00 表示。
请你返回矩阵中战斗力最弱的 kk 行的索引,按从最弱到最强排序。
如果第 ii 行的军人数量少于第 jj 行,或者两行军人数量相同但 ii 小于 jj,那么我们认为第 ii 行的战斗力比第 jj 行弱。
军人 总是 排在一行中的靠前位置,也就是说 11 总是出现在 00 之前。
示例 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
朴素解法
一个朴素的做法是对矩阵进行遍历,统计每一行的军人数量,并以二元组 (cnt_i, idx_i)(cnti,idxi) 的形式进行存储。
然后对所有行的战力进行排序,选出战力最小的 kk 个下标即是答案。
代码:
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; } } 复制代码
- 时间复杂度:遍历矩阵的复杂度为 O(m * n)O(m∗n);排序复杂度为 O(m\log{m})O(mlogm);构造答案复杂度为 O(k)O(k)。整体复杂度为 O(\max(m * n, m\log{m}))O(max(m∗n,mlogm))
- 空间复杂度:O(m)O(m) 空间用于存储所有的行战力;O(\log{m})O(logm) 空间用于排序。整体复杂度为 O(m + \log{m})O(m+logm)
二分 + 优先队列(堆)
根据「军人总是排在靠前位置」的特性,我们可以通过「二分」找到每一行最后一个军人的位置,从而在 O(\log{n})O(logn) 的复杂度内统计出每行的军人数量(战力情况)。
同时由于我们只需要「战力最弱」的 kk 行数据,这引导我们可以建立一个「战力大根堆」来做,「战力大根堆」存放着数量最多为 kk 的战力二元组 (cnt_i, idx_i)(cnti,idxi),堆顶元素为战力最大的数对。
每次通过「二分」得到当前行的战力值后,判断当前战力值与堆顶元素的战力大小关系:
- 如果当前战力值比堆顶的元素要大:直接丢弃当前战力值(不可能属于在第 kk 小的集合中);
- 如果当前战力值比堆顶的元素要小:将堆顶元素弹出,将当前行放入堆中。
一些细节:每次写二分都有同学会问,
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; } } 复制代码
- 时间复杂度:二分得到每行的战力情况,复杂度为 O(m\log{n})O(mlogn);堆中最多有 kk 个元素,将行信息存入堆中复杂度为 O(m\log{k});O(mlogk);构造答案复杂度为 O(k\log{k})O(klogk)。整体复杂度为 O(m * (\log{n} + \log{k}))O(m∗(logn+logk))
- 空间复杂度:O(k)O(k)
最后
这是我们「刷穿 LeetCode」系列文章的第 No.1337
篇,系列开始于 2021/01/01,截止于起始日 LeetCode 上共有 1916 道题目,部分是有锁题,我们将先把所有不带锁的题目刷完。
在这个系列文章里面,除了讲解解题思路以外,还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。
为了方便各位同学能够电脑上进行调试和提交代码,我建立了相关的仓库:github.com/SharingSour… 。
在仓库地址里,你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。