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

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

You are given an m x n binary matrix mat of 1's (representing soldiers) and 0's (representing civilians). The soldiers are positioned in front of the civilians. That is, all the 1's will appear to the left of all the 0's in each row.

A row i is weaker than a row j if one of the following is true:

The number of soldiers in row i is less than the number of soldiers in row j.

Both rows have the same number of soldiers and i < j.

Return the indices of the k weakest rows in the matrix ordered from weakest to strongest.

---

접근 방법: heap

• 주어진 매트릭스에서 0 앞에 오는 연속된 1값의 길이를 구한 후, 가장 짧은 순서대로 k개의 인덱스를 반환하되 값이 같으면 더 빠른 행을 우선으로 반환하는 문제이다.

• 각 행의 값을 모두 더한 힙 정렬을 이용하여 정렬시키고 k개의 요소를 뽑아내어 반환하는 식으로 코드를 작성했다.

• 단, 값이 같은 경우 더 앞에 오는 행의 인덱스를 반환하는게 문제의 조건이었으므로 처음 힙에 넣을 때 sum() 값과 행의 인덱스를 튜플로 묶어서 넣어주었다.

• heapq 모듈에서 값으로 튜플이 들어오면, 튜플의 가장 앞 요소부터 비교하는 식으로 대소비교를 해주므로, (val, index)의 형태로 넣어주면 앞의 값을 먼저 비교하여 작은 것을 앞에 놓고, 값이 같은 경우 뒤의 인덱스를 비교하여 작은 것을 앞에 놓게 되어 문제에서 요구하는 정답순으로 나열될 것이다.

코드

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        sums = [(sum(row), i) for i, row in enumerate(mat)]
        return [i for val, i in heapq.nsmallest(k, sums)]

---

구현

• 파이썬에서 제공하는 모듈을 사용하지 않고, 힙을 직접 구현하여 문제를 풀어보기로 하였다.

• 최소 힙에 필수적으로 필요한 요소는 다음과 같다.

  • heappush(item) : 힙에 요소를 추가
  • heappop() : 힙에서 가장 작은 값(root)을 뽑아 출력
  • percolate() : 힙의 조건에 맞게 요소들의 위치를 조정

• 요소를 추가한 후 요소의 위치를 위로 올라가면서 조절하는 percolate_up() 과, 요소를 제거한 후 가장 마지막 요소를 맨 앞으로 옮겨 내려가면서 조절하는 percolate_down() 을 따로 만들어서 구현하였다.

• 추가로, heappush 를 할 때, 요소로 무엇이 들어오냐에 따라 다른 방식으로 대소 비교를 해야하므로 is_smaller() 메소드를 만들어 대소 비교를 하였다.
  파이썬에서 부등호로 대소비교를 할 때, 튜플이 들어오면 자체적으로 내가 구현했던 방식으로 처리를 하여 결과를 출력하는 것을 확인할 수 있었다. 따라서 이 부분은 굳이 추가하지 않아도 제대로 동작하기 때문에 삭제하기로 하였다.

#    def _is_smaller(self, item1, item2):
#        if type(item1) == type(item2) == int:
#            return item1 < item2
#        elif type(item1) == type(item2) == tuple:
#            for p, q in zip(item1, item2):
#                if p == q:
#                    continue
#                else:
#                    return p < q
#        return False

• 그 후, 구현한 힙 클래스에 맞게 문제를 푸는 코드도 수정하였다.

from structure import BinaryHeap
from typing import List

class KWeakestRows:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        sums = [(sum(row), i) for i, row in enumerate(mat)]
        heap = BinaryHeap()

        for s in sums:
            heap.heappush(s)
        weak = [heap.heappop() for _ in range(k)]
        return [i for val, i in weak]