给出二维0-1矩阵,每行都由连续的1之后全0,输出含1最少的k行,二分确定每行0-1分隔的位置,heapq存每行的1的数量
1 #1337
2 #Runtime: 71 ms (Beats 88.26%)
3 #Memory: 13.6 MB (Beats 61.50%)
4
5 class Solution(object):
6 def kWeakestRows(self, mat, k):
7 """
8 :type mat: List[List[int]]
9 :type k: int
10 :rtype: List[int]
11 """
12 def bi_sect(x):
13 l, r = 0, len(mat[0])
14 while l < r:
15 mid = (l + r) // 2
16 if mat[x][mid] == 1:
17 l = mid + 1
18 else:
19 r = mid
20 return l
21
22 hq = []
23 ans = []
24 for r in xrange(len(mat)):
25 p = bi_sect(r)
26 heapq.heappush(hq, (p, r))
27 for _ in xrange(k):
28 ans.append(heappop(hq)[1])
29 return ans