퀵 정렬 - (Quick-Sort) with Python

📖 문제 설명

문제 : https://www.acmicpc.net/problem/24090

• 오늘도 서준이는 퀵 정렬 수업 조교를 하고 있다. 아빠가 수업한 내용을 학생들이 잘 이해했는지 문제를 통해서 확인해보자.
• N개의 서로 다른 양의 정수가 저장된 배열 A가 있다. 퀵 정렬로 배열 A를 오름차순 정렬할 경우 배열 A에 K 번째 교환되는 수를 구해서 우리 서준이를 도와주자.
• 크기가 N인 배열에 대한 퀵 정렬 의사 코드는 다음과 같다.

quick_sort(A[p..r]) { # A[p..r]을 오름차순 정렬한다.
    if (p < r) then {
        q <- partition(A, p, r);  # 분할
        quick_sort(A, p, q - 1);  # 왼쪽 부분 배열 정렬
        quick_sort(A, q + 1, r);  # 오른쪽 부분 배열 정렬
    }
}

partition(A[], p, r) {
    x <- A[r];    # 기준원소
    i <- p - 1;   # i는 x보다 작거나 작은 원소들의 끝지점
    for j <- p to r - 1  # j는 아직 정해지지 않은 원소들의 시작 지점
        if (A[j] ≤ x) then A[++i] <-> A[j]; # i값 증가 후 A[i] <-> A[j] 교환
    if (i + 1 != r) then A[i + 1] <-> A[r]; # i + 1과 r이 서로 다르면 A[i + 1]과 A[r]을 교환
    return i + 1;
}

💡요구사항 분석

퀵정렬

• 리스트 가운데서 하나의 원소를 고른다. 이렇게 고른 원소를 피벗이라고 합니다.
• 피벗 앞에는 피벗보다 값이 작은 모든 원소들이 오고, 피벗 뒤에는 피벗보다 값이 큰 모든 원소들이 오도록 피벗을 기준으로 리스트를 둘로 나눈다. 이렇게 리스트를 둘로 나누는 것을 분할이라고 한다. 분할을 마친 뒤에 피벗은 더 이상 움직이지 않습니다.
• 분할된 두 개의 작은 리스트에 대해 재귀적으로 이 과정을 반복한다. 재귀는 리스트의 크기가 0이나 1이 될 때까지 반복됩니다.

코드

import sys

sys.setrecursionlimit(10000)

n, k = map(int, sys.stdin.readline().split())
arr = list(map(int, sys.stdin.readline().split()))

def partition(arr, start, end):
    pivot = arr[end]  # 피복은 배열 마지막으로 설정
    i = start - 1  # 작은 값들을 배치할 인덱스

    for j in range(start, end):
        if arr[j] <= pivot:
            i += 1
            arr[i], arr[j] = arr[j], arr[i]

    arr[i + 1], arr[end] = arr[end], arr[i + 1]  # 피봇 갱신
    return i + 1

def quick_sort(arr, start, end):
    if start < end:
        mid = partition(arr, start, end)
        quick_sort(arr, start, mid - 1) // 왼쪽
        quick_sort(arr, mid + 1, end) // 오른쪽 

quick_sort(arr, 0, len(arr) - 1)

print(arr)

📖 코드 풀이

import sys
sys.setrecursionlimit(10000)

n, k = map(int, sys.stdin.readline().split())
arr = list(map(int, sys.stdin.readline().split()))
answer = []

def partition(arr, start, end):
    global k
    pivot = arr[end]
    i = start - 1

    for j in range(start, end):
        if arr[j] > pivot:
            i += 1
            arr[i], arr[j] = arr[j], arr[i]

            k -= 1
            if k == 0:
                answer.append(arr[i])
                answer.append(arr[j])
                answer.sort()
                print(answer[0], answer[1])
                exit()

    if i + 1 != end:
        arr[i + 1], arr[end] = arr[end], arr[i + 1]

        k -= 1
        if k == 0:
            answer.append(arr[i + 1])
            answer.append(arr[end])
            answer.sort()
            print(answer[0], answer[1])
            exit()

    return i + 1

def quick_sort(arr, start, end):
    if start < end:
        mid = partition(arr, start, end)
        quick_sort(arr, start, mid - 1)
        quick_sort(arr, mid + 1, end)

quick_sort(arr, 0, len(arr) - 1)

print(-1)

• 재귀제한 때문에 런타임에러가 발생하여 sys.setrecursionlimit(10000) 설정해줘야 합니다.
• 퀵정렬은 최악의 경우 N^2 이기 때문에 값은 찾았다면 exit() 을 통해 프로그램을 종료해줍니다.