[리트코드] 1337. 행렬의 K개 가장 약한 행 (Python)

The K Weakest Rows in a Matrix - LeetCode

문제

m x n이진 행렬 mat이 제공 되는데, 1의 개수가 가장 적은 k 개의 인덱스를 반환하는 프로그램을 작성하시오.

 

과정

1. 힙을 활용해 문제를 풀어준다. heapq는 최소힙이다.
2. 입력 받은 각 리스트들이 1을 가진 갯수를 가지고 heappush 해준다.
3. 그 후 heappop을 통해 가장 작은 값부터 정답 리스트에 넣어준다.

 

입력

[[1,1,0,0,0],[1,1,1,1,0],[1,0,0,0,0],[1,1,0,0,0],[1,1,1,1,1]]
3

출력

[2,0,3]

 

풀이

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        heap = []
        answer = []
        for i in range(len(mat)):
            cnt = mat[i].count(1)
            heapq.heappush(heap, [cnt, i])
        for _ in range(k):
            answer.append(heapq.heappop(heap)[1])

        return answer

 

한 줄씩 코드 해석해보기

heap = []
answer = []

• 힙으로 바꿔줄 리스트와 정답을 리턴 할 리스트를 만들어준다.

for i in range(len(mat)):
    cnt = mat[i].count(1)
    heapq.heappush(heap, [cnt, i])

• 입력받은 리스트안의 각 리스트들을 돌면서 1의 개수를 센 후 heap 리스트에 heappush 해준다.
• 이때 개수와 인덱스가 같이 push된다.
• heapq는 최소힙의 기능을 갖고 있기 때문에 작은 값부터 순서대로 push 된다.

for _ in range(k):
    answer.append(heapq.heappop(heap)[1])

• pop하면 우선순위가 가장 작은 값부터 answer에 append된다.
• 인덱스의 값만 필요하기 때문에 [1] 위치만 pop해준다.