Algorithm/릿코드

99클럽 코테 스터디 29일차 1337. The K Weakest Rows in a Matrix

Albosa2lol 2024. 6. 27. 22:56

 

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

 

주어진 m x n 크기의 이진 행렬 mat는 1(병사)와 0(민간인)으로 이루어져 있습니다. 이 행렬에서 병사들은 민간인 앞에 배치되어 있습니다. 즉, 각 행에서 모든 1은 모든 0의 왼쪽에 나타납니다.

행 i가 행 j보다 약하다는 것은 다음 두 가지 조건 중 하나에 해당합니다:

  1. 행 i의 병사 수가 행 j의 병사 수보다 적다.
  2. 두 행의 병사 수가 같고, i가 j보다 작다.

가장 약한 행부터 가장 강한 행까지의 순서로 k개의 가장 약한 행의 인덱스를 반환하세요.

 

 

 

 

 

 

 

풀이

 

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        # 1. 각 행의 병사 수를 센다.
        soldier_counts = []
        for i in range(len(mat)):
            count = 0
            for j in range(len(mat[i])):
                if mat[i][j] == 1:
                    count += 1
                else:
                    break  # 행렬의 조건에 따라 1 이후는 모두 0이므로 더 이상 셀 필요가 없음
            soldier_counts.append((i, count))
        
        # 2. 병사 수를 기준으로 정렬 (병사 수가 같으면 행 번호로 정렬)
        sorted_rows = sorted(soldier_counts, key=lambda x: (x[1], x[0]))
        
        # 3. 가장 약한 k개의 행의 인덱스를 반환
        weakest_rows = []
        for i in range(k):
            weakest_rows.append(sorted_rows[i][0])
        
        return weakest_rows