m1ndy5's coding blog

LeetCode 1337. The K Weakest Rows in a Matrix with Python 본문

알고리즘 with python/알고리즘 스터디

LeetCode 1337. The K Weakest Rows in a Matrix with Python

정민됴 2024. 1. 5. 10:44

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

import heapq

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        order = [row.count(1) for row in mat]
        heap = order[:]
        heapq.heapify(heap)
        visited = [0 for _ in range(len(order))]
        answer = []
        for _ in range(k):
            m = heapq.heappop(heap)
            for i, v in enumerate(order):
                if v == m and visited[i] == 0:
                    answer.append(i)
                    visited[i] = 1
                    break

        return answer

1의 개수로 어떤 행이 제일 약한지 알아내야하는 문제였다.
만약 1의 개수가 같다면 행 번호가 더 작은 애가 약한 것이었기 때문에 반복문을 돌려서 해결하였다.