题目描述

给你一个大小为 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]

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/the-k-weakest-rows-in-a-matrix
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

思路分析

• 这是一个简单题目，首先需要理解题意，题意简述为，统计每一个子数组中 1 的个数，根据 1 的个数对二维数组索引进行排序，当 1 的个数相同的时候，按照自然升序排序。
• 首先使用朴素解法，下面代码使用 TreeMap 数据结构，实现了排序功能。TreeMap 实现了红黑树的结构,是一个能比较元素大小的Map集合，会对传入的key进行了大小排序。TreeMap 可以使用元素的自然顺序，也可以使用集合中自定义的比较器来进行排序。

代码

class Solution {
    public int[] kWeakestRows(int[][] mat, int k) {
        int[] ans = new int[k];
        int n = mat.length;
        int[] nums = new int[n];
        Map<Integer, List<Integer>> map = new TreeMap<>();
        for (int i = 0; i < n; i++) {
            int cnt = 0;
            for (int j = 0; j < mat[i].length; j++) {
                if (mat[i][j] == 1) {
                    cnt++;
                } else {
                    break;
                }
            }
            List<Integer> cntSameKeys = new ArrayList<>();
            if (map.containsKey(cnt)) {
                cntSameKeys = map.get(cnt);
            } 
            cntSameKeys.add(i);
            map.put(cnt, cntSameKeys);
        }

        int i = 0;
        while (i < k) {
            for (Integer num : map.keySet()) {
                List<Integer> keys = map.get(num);
                keys.sort(new Comparator<Integer>() {
                    @Override
                    public int compare(Integer o1, Integer o2) {
                        return o1 - o2;
                    }
                });
                
                for (Integer idx : keys) {
                    if (i >= k) {
                        break;
                    }
                    ans[i] = idx;
                    i++;
                }
            }
        }
        return ans;
    }
}

总结

• 上述代码的时间复杂度是 O(n * m), 空间复杂度是 O(n)
• 坚持算法每日一题，夯实基础，加油！