N개의 수가 주어지고 이를 오름차순으로 정렬하였을 때 앞에서부터 K번째 있는 수를 어떻게 구할 것인지 구하는 문제이다.
현재 데이터 크기가 최대 5,000,000이고 시간제한은 2초이므로 의 시간복잡도를 가지는 알고리즘을 사용해야하므로 병합정렬을 사용하는 것이 안전하다.
하지만 필자는 이번 시간에 퀵 정렬에 대한 원리를 파악하고, 이를 활용한 퀵 셀렉트 알고리즘을 통해 풀이를 진행하니 참고하길 바란다.
(퀵 정렬은 pivot 선정에 따라 최악의 경우 )
먼저 퀵 정렬이란 기준값인 pivot을 선정하여 해당 값보다 작은 데이터와 큰 데이터로 분류하는 것을 반복하여 정렬하는 알고리즘이다.
알고리즘 흐름은 다음과 같다.
-
데이터를 분할하는 pivot 설정
=> 설정하는 기준은 가장 오른쪽, 가장 왼쪽, 또는 중간 위치 등 문제 요건에 맞춰서 자유롭게 설정하면 됨 -
pivot을 기준으로 다음 2-1 ~ 2-5 과정을 거쳐 데이터를 2개의 집합으로 분리한다.
2-1. start을 가리키는 데이터가 pivot이 가리키는 데이터보다 작으면 start를 오른쪽으로 1칸 이동
2-2. end가 가리키는 데이터가 pivot이 가리키는 데이터보다 크면 end를 왼쪽으로 1칸 이동
2-3. start가 가리키는 데이터가 pivot이 가리키는 데이터보다 크고, end가 가리키는 데이터가 pivot이 가리키는 데이터보다 작을경우 start와 end가 가리키는 데이터를 swap하고 start는 오른쪽, end는 왼쪽으로 1칸씩 이동
2-4. start와 end가 서로 교차할때까지 2-1 ~ 2-3 과정 반복
2-5. start와 end가 교차한 지점에서 가리키는 데이터와 pivot이 가리키는 데이터를 swap한다.
-
분리 집합에서 각각 다시 pivot을 선정
=> 이 과정에서 함수를 재귀호출하게 됨 -
분리 집합이 1개 이하가 될때까지 1~3 과정을 반복
그렇다면 위 알고리즘을 어떻게 적용해야 K번째 수를 효율적으로 찾아낼 수 있을까?
기준값이 pivot이 K번째 수라면 바로 알고리즘을 종료하고 정답을 출력하면 되지 않을까?
따라서 문제풀이 방법은 다음과 같다.
quickSort(start, end, K) 함수 선언
- K번째 값이 pivot 위치인 경우 정답 출력
- K번째 값이 pivot보다 오른쪽 구간에 있다면 오른쪽만 탐색
=> quickSort(pivot+1, end, K) - K번째 값이 pivot보다 왼쪽 구간에 있다면 왼쪽만 탐색
=> quickSort(start, pivot-1, K)
findPivot(start, end) 함수 선언
-
만약 데이터가 2개인 경우 바로 비교한 후 원소 정렬 and 피봇의 위치 return
-
중간 위치를 pivot으로 설정한 다음 맨 앞에 있는 데이터와 swap
=> start(i 인덱스), end(j 인덱스) 이동을 편리하게 진행하기 위함 -
j가 pivot보다 큰 경우 j-- 연산을 반복
=> 오른쪽에서부터 pivot보다 작은 값이 있을때까지 j-- 연산 반복한다는 말과 동일 -
i가 pivot보다 작으면서 i가 j보다 작은 경우 i++ 연산을 반복
=> 왼쪽에서부터 pivot보다 큰 값이 있을때까지 i++ 연산 반복한다는 말과 동일 -
i와 j가 서로 교차하지 못한 채 i++ or j-- 연산이 종료되었다면 서로 가리키는 데이터를 swap한 후 i++, j-- 연산 진행
-
3 ~ 5번의 반복문이 끝나면 i, j가 서로 교차했다는 의미이므로, i와 j가 만난 위치가 가리키는 데이터와 pivot이 가리키는 데이터를 swap
=> i와 j 중 더 작은 값을 가리키는 데이터와 pivot 값을 swap한 후 pivot의 최종 위치를 반환 -
반환한 최종 위치가 K번째 수인지 비교하여 정답 출력 or 함수 재귀호출하여 선택한 그룹 안에서 다시 탐색 진행
이를 모두 코드에 녹여내면 아래와 같다...
import sys
input = sys.stdin.readline
N, K = map(int, input().split())
A = list(map(int, input().split()))
# 퀵 정렬 알고리즘
def quickSort(start, end, K):
global A
# 분할할 구간이 아직 남은경우
if start < end:
# 현재 구간의 피봇을 설정하고 배열을 분할
pivot = findPivot(start, end)
# 목표값인 K번째에 해당하는 값이 피봇 위치라면
if pivot == K:
return # 바로 종료하여 오름차순으로 정렬된 K번째 값 출력
# K번째 값이 pivot보다 오른쪽 구간에 있다면 오른쪽만 탐색
elif pivot < K:
quickSort(pivot + 1, end, K)
# K번째 값이 pivot보다 왼쪾 구간에 있다면 왼쪽만 탐색
elif pivot > K:
quickSort(start, pivot-1, K)
# 배열의 두 원소를 교환하는 함수
# 교환할 두 원소의 인덱스를 넣으면 됨
def swap(i, j):
global A
temp = A[i]
A[i] = A[j]
A[j] = temp
# e.g. [3, 7, 2, 5, 1, 6, 4]
# 피봇을 찾고 배열을 재배치하는 함수
def findPivot(start, end):
global A
# 데이터가 2개인 경우는 바로 비교하여 정렬
if start + 1 == end:
# 두 원소가 정렬되어 있지 않다면 교환
if A[start] > A[end]:
swap(start, end)
return end # 피봇의 위치 반환
# 중간 인덱스를 구하여 피봇으로 설정한 다음
M = (start + end) // 2
# 시작 위치와 교환
swap(start, M)
pivot = A[start] # 피봇값
i = start + 1 # 피봇의 다음 인덱스를 i로 설정
j = end # 배열의 마지막 인덱스를 j로 설정
# i와 j가 서로 교차할때까지 반복
while i <= j:
# 오른쪽에서부터 피봇보다 작은 값이 있을때까지 반복하여 내려옴
while pivot < A[j] and j > 0:
j = j - 1
# 왼쪽에서부터 피봇보다 큰 값이 있을때까지 반복하여 올라감
while pivot > A[i] and i < len(A)-1:
i = i + 1
# i와 j가 서로 교차하지 못한 채 반복문이 종료되었다면
# 1. M = pivot = 5, [5, 7, 2, 3, 1, 6, 4]
# 2. i는 현재 1번 인덱스인 7을 가리키고, j는 6번 인덱스인 4를 가리키는 상황
if i <= j:
# 서로 다른 집합에 있는 두 값을 swap
swap(i, j)
i = i + 1
j = j - 1
# 위의 while문이 끝나면 i와 j가 교차한 지점이 되므로
# 피봇이 가리키는 값과 i와 j가 교차한 지점과 swap
# 3. j = 4이고, 현재 1을 가리킴
# 4. i = 5이고, 현재 6을 가리킴
# 5. 따라서 pivot보다 작은 값을 가리키고 있는 A[j]와 pivot값을 swap
A[start] = A[j]
A[j] = pivot
# 피봇의 최종위치 반환
return j
# K번째 수를 찾기 위해 퀵 정렬 실행
quickSort(0, N - 1, K - 1) # 배열은 0부터 시작하니까 K-1을 넣어야 함
print(A[K - 1]) # 결과 출력 역시 K-1우리가 아는 퀵 정렬에서 K번째 수를 찾게되면 더 이상 정렬을 하지 않고 바로 종료하는 방식이다.
함수를 통해 구현해야하다보니 코드가 길어지므로 시간을 갖고 상세하게 분석해보길 바란다.
