파이썬에서 행렬의 각 열을 오름차순으로 정렬하는 방법

문제 소개

행렬(matrix)이 주어졌을 때, 각 열(column)의 요소들을 오름차순으로 정렬하는 문제입니다. 일반적인 정렬은 행 전체를 대상으로 하지만, 이 문제는 열 단위로 정렬해야 한다는 점이 특징입니다.

예를 들어 입력이 다음과 같다면,

출력은 다음과 같습니다.

첫 번째 열은 [11, 6, 1] → [1, 6, 11], 두 번째 열은 [21, 6, 11] → [6, 11, 21], 세 번째 열은 [31, 4, 8] → [4, 8, 31]로 각각 정렬된 것을 확인할 수 있습니다.

해결 알고리즘

다음 단계에 따라 문제를 해결할 수 있습니다.

• R := 행렬의 행 개수, C := 행렬의 열 개수로 설정합니다.
• res := 주어진 행렬과 같은 크기의, 0으로 초기화된 결과 행렬을 생성합니다.
• col을 0부터 C-1까지 반복합니다.
  • values := 각 행에서 col번째 요소들을 모아 만든 리스트(열 벡터)를 구합니다.
  • row를 0부터 R-1까지 반복합니다.
    • res[row][col] := values에서 마지막 요소를 꺼내 저장합니다.
• res를 반환합니다.

여기서 values를 내림차순으로 정렬한 뒤 뒤에서부터 하나씩 꺼내면(pop), 결과적으로 오름차순으로 채워지게 됩니다.

파이썬 구현 예제

class Solution:
    def solve(self, matrix):
        R = len(matrix)
        C = len(matrix[0])
        res = [[0] * C for _ in range(R)]
        for col in range(C):
            values = [r[col] for r in matrix]
            values.sort(reverse=True)
            for row in range(R):
                res[row][col] = values.pop()
        return res

ob = Solution()
matrix = [[11, 21, 31], [6, 6, 4], [1, 11, 8]]
print(ob.solve(matrix))

입력

[[11, 21, 31],
[6, 6, 4],
[1, 11, 8]]

출력

[[1, 6, 4], [6, 11, 8], [11, 21, 31]]

코드 동작 원리

핵심 아이디어를 정리하면 다음과 같습니다.

• 열 추출: 리스트 컴프리헨션 [r[col] for r in matrix]를 사용해 특정 열의 값만 모읍니다.
• 내림차순 정렬 후 pop: sort(reverse=True)로 내림차순 정렬하면 리스트의 마지막 요소가 가장 작은 값이 됩니다. 따라서 pop()으로 뒤에서부터 꺼내면 작은 값부터 차례대로 저장되어 최종적으로 오름차순이 됩니다.
• 시간 복잡도: 각 열마다 O(R log R)의 정렬이 수행되므로, 전체 시간 복잡도는 O(C × R log R)입니다.

참고: zip을 활용한 더 간결한 방법

파이썬에서는 zip(*matrix)로 행렬을 전치(transpose)하면 훨씬 간결하게 해결할 수 있습니다.

def solve(matrix):
    return [list(col) for col in zip(*[sorted(r) for r in zip(*matrix)])]

먼저 zip(*matrix)로 행렬을 전치해 열을 행처럼 다룬 뒤 각 행(원래의 열)을 정렬하고, 다시 한 번 전치하여 원래 형태로 되돌리는 방식입니다. 파이썬다운 관용적인 표현(idiomatic way)으로 짧은 코드로 같은 결과를 얻을 수 있습니다.