Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 정렬된 행렬에서 n번째로 작은 수 찾기

각 행과 열이 비내림차순(오름차순)으로 정렬되어 있는 2차원 행렬이 주어졌을 때, 이 행렬에서 n번째로 작은 수를 찾는 문제입니다.

예를 들어 다음과 같은 행렬이 입력으로 주어진다고 가정해 보겠습니다.

2430
3431
6632

이때 n = 4라면, 출력 결과는 6이 됩니다.

문제 해결 접근 방법

이 문제는 다음 단계를 따라 간단하게 해결할 수 있습니다.

  • 빈 리스트(lst)를 하나 생성합니다.
  • 행렬의 각 행 i에 대해 반복합니다.
    • 행 i의 각 요소 j에 대해 반복하며, j를 lst의 끝에 추가합니다.
  • 리스트 lst 전체를 오름차순으로 정렬합니다.
  • lst[n] 값을 반환합니다.

즉, 2차원 행렬의 모든 원소를 1차원 리스트로 펼친 뒤 정렬하고, 인덱스 n에 해당하는 값을 꺼내면 됩니다. 아래 구현 예제를 통해 더 자세히 이해해 보겠습니다.

예제 코드

class Solution:
   def solve(self, matrix, n):
      lst = []
      for i in matrix:
         for j in i:
            lst.append(j)
      lst.sort()
      return lst[n]
ob = Solution()
matrix = [ [2, 4, 30], [3, 4, 31], [6, 6, 32] ]
n = 4
print(ob.solve(matrix, n))

입력

matrix = [
[2, 4, 30],
[3, 4, 31],
[6, 6, 32] ]
n = 4

출력

6

복잡도 분석

위 알고리즘의 시간 복잡도와 공간 복잡도는 다음과 같습니다.

  • 시간 복잡도: O(m × k log(m × k)) — m은 행의 개수, k는 열의 개수입니다. 모든 원소를 모은 뒤 정렬하기 때문에 정렬 비용이 지배적입니다.
  • 공간 복잡도: O(m × k) — 행렬의 모든 원소를 저장할 1차원 리스트가 필요합니다.

참고로, 행렬의 크기가 매우 큰 경우에는 힙(heap) 자료구조나 이분 탐색을 활용하면 더 효율적으로 n번째 작은 수를 찾을 수 있습니다. 하지만 위 방법은 구현이 간단하고 직관적이어서 대부분의 일반적인 상황에서 충분히 실용적입니다.