n×n 크기의 정수 행렬과 값 k가 주어졌을 때, 2차원(2D) 배열에서 k번째로 작은 요소를 찾는 것이 이 글의 목표입니다. 이 문제는 파이썬의 heapq 모듈을 활용하면 아주 간단하게 해결할 수 있습니다.
파이썬에서 힙 큐(Heap Queue)는 heapq 모듈을 통해 제공됩니다. 이 모듈은 최소 힙(min heap) 구조를 기반으로 동작하며, 힙에서 요소를 꺼낼(pop) 때마다 항상 가장 작은 값이 먼저 반환됩니다. 특히 nsmallest() 메서드를 사용하면 데이터 집합에서 가장 작은 n개의 값을 한 번에 추출할 수 있어, k번째 작은 요소를 구하는 데 매우 유용합니다.
예시
입력 배열:: 10 20 20 40 15 45 40 30 32 33 30 50 12 78 99 78 k 값은 10 → 10번째로 작은 요소는 40
알고리즘
- 먼저 2D 배열을 생성합니다.
- 첫 번째 행을 변수에 할당한 뒤
heapify()를 이용해 최소 힙으로 변환합니다. - 나머지 행들을 순회하면서 모든 요소를
heappush()로 최소 힙에 삽입합니다. nsmallest(k, iterable)메서드로 가장 작은 k개의 요소 리스트를 구한 뒤, 그 리스트의 마지막 요소를 출력합니다.
예제 코드
# 파이썬으로 2D 배열에서 K번째로 작은 요소를 찾는 프로그램
import heapq
def smallestele(A):
assignval = A[0]
heapq.heapify(assignval)
for i in A[1:]:
for j in i:
heapq.heappush(assignval, j)
mini = heapq.nsmallest(k, assignval)
print(k, "번째로 작은 요소는", mini[-1])
# 드라이버 프로그램
if __name__ == "__main__":
A = []
n = int(input("N x N 행렬의 N 값 입력 : ")) # 예: 3
# 2D 배열을 저장하기 위한 리스트
# 사용자 입력을 받아 리스트에 저장
print("요소를 입력하세요 ::>")
for i in range(n):
row = [] # 행을 임시 저장하는 리스트
for j in range(n):
row.append(int(input())) # 입력값을 행 리스트에 추가
A.append(row) # 완성된 행을 배열에 추가
print(A)
# [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
# 2D 배열을 행렬 형태로 출력
print("행렬 형태로 배열 출력")
for i in range(n):
for j in range(n):
print(A[i][j], end=" ")
print() # 줄바꿈
k = int(input("k번째 위치 입력 ::>"))
smallestele(A)
실행 결과
N x N 행렬의 N 값 입력 : 4 요소를 입력하세요 ::> 10 20 20 40 15 45 40 30 32 33 30 50 12 78 99 78 [[10, 20, 20, 40], [15, 45, 40, 30], [32, 33, 30, 50], [12, 78, 99, 78]] 행렬 형태로 배열 출력 10 20 20 40 15 45 40 30 32 33 30 50 12 78 99 78 k번째 위치 입력 ::>10 10 번째로 작은 요소는 40
마무리
이 방식은 첫 번째 행을 힙으로 만든 후 나머지 모든 요소를 힙에 삽입하고, nsmallest()로 가장 작은 k개를 추출하는 구조입니다. 요소의 총 개수를 N=n²이라 할 때 시간 복잡도는 약 O(N log N) 수준이며, 구현이 단순하고 직관적이라는 장점이 있습니다. 참고로 행렬의 각 행과 열이 이미 오름차순으로 정렬되어 있는 특수한 경우에는 힙을 이용한 병합 기법 등 더 효율적인 접근 방법도 존재합니다.