1337. The K Weakest Rows in a Matrix

Runtime    2 ms
Beats      80.96%


Memory     44.3 MB
Beats      86.29%


class Solution {
    public int[] kWeakestRows(int[][] mat, int k) {
        PriorityQueue <int[]> ans=new PriorityQueue<>((a,b)->a[0]!=b[0]?b[0]-a[0]:b[1]-a[1]);

        int res[]=new int[k];

        for(int i=0;i<mat.length;i++){
            int count=0;
            for(;count<mat[i].length;count++){
                if(mat[i][count] == 0)
                    break;
            }
            ans.offer(new int[]{count,i});
        }

        while(ans.size()>k){
            ans.poll();
        }

        while(k>0){
            res[--k]=ans.poll()[1];
        }

        return res;
    }
}