알고리즘/파이썬

[파이썬 알고리즘] 1337. leetcode- The K Weakest Rows in a Matrix

귤먹는코더 2022. 5. 25. 15:56
728x90

문제 설명 

1은 군인, 0은 민간인이다. 각 열에서 제일 군인이 적은 열을 찾고 k 만큼 반환해라

 

import heapq
from typing import List


class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        power = []
        for r, row in enumerate(mat):
            tmp = 0
            for n in row:
                if n:
                    tmp += 1
                else:
                    break
            power.append((tmp, r))
        heapq.heapify(power)
        return [heapq.heappop(power)[1] for _ in range(k)]
728x90