【Leetcode】1337. The K Weakest Rows in a Matrix

题目地址：

https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/

给定一个$m$行$n$列的矩阵，每行只有$0$和$1$，并且任意一个$1$都在任意一个$0$的左边。再给定一个正整数$k$，要求返回前$k$个$1$的个数最小的行的下标，如果有并列则按行下标从小到大排列。

代码如下：

import java.util.Arrays;

public class Solution {
    public int[] kWeakestRows(int[][] mat, int k) {
        int[][] info = new int[mat.length][];
        for (int i = 0; i < mat.length; i++) {
        	// 第一维是下标，第二维是1的个数，1的个数可以用二分来求
            info[i] = new int[]{i, cnt(mat[i])};
        }
        
        Arrays.sort(info, (r1, r2) -> r1[1] != r2[1] ? Integer.compare(r1[1], r2[1]) : Integer.compare(r1[0], r2[0]));
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = info[i][0];
        }
        
        return res;
    }
    
    private int cnt(int[] A) {
        int l = 0, r = A.length - 1;
        while (l < r) {
            int m = l + (r - l + 1 >> 1);
            if (A[m] == 1) {
                l = m;
            } else {
                r = m - 1;
            }
        }
        
        return A[l] == 1 ? l + 1 : 0;
    }
}

时间复杂度$O(m\log n+m\log m)$，空间$O(m)$。