[Leetcode] 1337. The K Weakest Rows in a Matrix

✅ 문제

 

 

 

 

---

 

✅ 나의 풀이

import heapq
class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        
        # 리스트에 1이 나온 개수를 카운트해서 순서대로 담음
        min_heap = [(row.count(1), i) for i, row in enumerate(mat)]
            
        # 리스트 -> 힙
        heapq.heapify(min_heap)
        
        # 답을 반환할 리스트 선언
        answer = []

        for i in range(k):
            answer.append(heapq.heappop(min_heap)[1])

        return answer

 

collections.Counter를 이용할까 하다가 row.count를 이용해보았다.

mat 리스트에 담긴 각 리스트마다 1의 개수를 세어 index를 붙이고 최소힙으로 만드는 방법이다.

최소힙으로 만들게 되면 

min_heap = [(1, 2), (2, 3), (2, 0), (4, 1), (5, 4)] 

이런식으로 각 리스트의 앞자리를 기준으로 정렬이 되는데,

heappop을 하게 됐을 때 비교하는 부모노드가 같은 수일 경우 자식 노드를 비교하여 크기가 작은 것부터 나오게 된다.

따라서 (1,2) (2,0), (2,3) 이 뽑혀 [2, 0, 3]이 출력되는 것을 확인할 수 있다.