문제 링크

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

개요

• 2차원 배열에 대해 각각의 리스트의 합을 기준으로 정렬을 하고 그 순서를 반환하는 문제이다.
• 파이썬에서는 내장함수 sort()를 사용하면 쉽게 풀 수 있다.

문제 해설

• 입력으로 2차원 배열 mat과 출력값의 개수를 의미하는 정수 k가 주어진다.
• mat에 있는 각각의 리스트는 0과 1의 조합으로 이루어져 있으며 1의 개수가 많은 리스트가 강한 리스트이다.
• 문제에서 요구하는 것은 1. 리스트를 약한 순으로 정렬하고
  2. 정렬하기 전의 인덱스 번호를 정렬된 순서대로 반환하는 것이다.
• 이를 위해 리스트의 인덱스 번호와 리스트의 합을 따로 저장할 필요가 있으므로 for문을 통해 mat을 순회한다.
• 순회하면서 mat의 각 리스트 내용을 [인덱스 번호, 리스트의 합]으로 덮어쓰고
  이후에 1번 원소(리스트의 합)을 기준으로 mat를 정렬한다.
• 마지막에 정렬된 mat의 0번 원소(인덱스 번호)를 k개 만큼만 추출해서 반환하면 된다.

해설 코드

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class Solution(object):
    def kWeakestRows(self, mat, k):
        """
        :type mat: List[List[int]]
        :type k: int
        :rtype: List[int]
        """

        for i in range(len(mat)):
            mat[i] = [i, sum(mat[i])]

        mat.sort(key=lambda x: x[1])

        return [mat[i][0] for i in range(k)]