-
Notifications
You must be signed in to change notification settings - Fork 1
/
solution.py
31 lines (23 loc) · 860 Bytes
/
solution.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution:
def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
heap = []
def countSoldier(row):
left, right = 0, len(row)
while left < right:
mid = left + (right - left) // 2
if row[mid] == 1:
left = mid + 1
else:
right = mid
return left
for index, row in enumerate(mat):
curr = (-countSoldier(row), -index)
if len(heap) == k:
heapq.heappushpop(heap, curr)
else:
heapq.heappush(heap, curr)
res = []
for _ in range(k):
_, index = heapq.heappop(heap)
res.append(-index)
return res[::-1]