[리트코드(LeetCode)] 1337번 The K Weakest Rows in a Matrix

문제

https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/

The K Weakest Rows in a Matrix - LeetCode

행렬에서 가장 약한 행의 인덱스를 약한 순서대로 k번째까지 반환하는 문제다. (1은 병사, 0은 민간인을 나타냄)

예시

Example 1:

Input: mat = [[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]],
	k = 3
Output: [2,0,3]
Explanation: 
The number of soldiers in each row is: 
- Row 0: 2 
- Row 1: 4 
- Row 2: 1 
- Row 3: 2 
- Row 4: 5 
The rows ordered from weakest to strongest are [2,0,3,1,4].

 

Example 2:

Input: mat = [[1,0,0,0], [1,1,1,1], [1,0,0,0], [1,0,0,0]], 
	k = 2
Output: [0,2]
Explanation: 
The number of soldiers in each row is: 
- Row 0: 1 
- Row 1: 4 
- Row 2: 1 
- Row 3: 1 
The rows ordered from weakest to strongest are [0,2,3,1].

풀이

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        heap = []
        
        for i, row in enumerate(mat):
            strength = row.count(1)
            heappush(heap, (strength, i))

        res = []
        for _ in range(k):
            strength, idx = heappop(heap)
            res.append(idx)

        return res

 

- 각 행의 취약성을 계산하고 순서대로 반환하기 위해 heap 자료구조를 사용한다.

- 취약성이 가장 낮은 행이 루트에 위치하도록 병사들의 합(힘..?)과 인덱스를 튜플로 힙에 넣어준다.

- 가장 낮은 행부터 k번만큼 차례대로 꺼내 해당하는 인덱스를 결과 리스트에 추가하고 반환한다.