[백준.11004 - Python] K번째 수

📌 11004. K번째 수

수 N개 A1, A2, ..., AN이 주어진다. A를 오름차순 정렬했을 때, 앞에서부터 K번째 있는 수를 구하는 프로그램을 작성하시오

 

 

✏️ 입력. 

첫째 줄에 N(1 ≤ N ≤ 5,000,000)과 K (1 ≤ K ≤ N)이 주어진다.

둘째에는 A1, A2, ..., AN이 주어진다. (-109 ≤ Ai ≤ 109)

 

📋 출력. 

A를 정렬했을 때, 앞에서부터 K번째 있는 수를 출력한다.

 

 

import sys

n, m = map(int, sys.stdin.readline().split())
arr = list(map(int, sys.stdin.readline().split()))
arr.sort()
print(arr[m-1])

 

내장 함수를 사용하면 간단하게 풀 수 있는 문제다 

 

 

import sys

# 첫번째 데이터를 피벗으로 (호어 분할 방식)
def q_sort(arr,start,end):
    if start < end:
        pivot = arr[end]
        idx = start
        
        for i in range(start, end):
            if arr[i] < pivot:
                arr[idx], arr[i] = arr[i], arr[idx]
                idx += 1
        arr[idx], arr[end] = arr[end], arr[idx]
        q_sort(arr, start, idx-1)
        q_sort(arr, idx+1, end)
        
if __name__ == "__main__":
    n, m = map(int, sys.stdin.readline().split())
    arr = list(map(int, sys.stdin.readline().split()))
    q_sort(arr,0,n-1)    
    print(arr[m-1])

 

퀵 정렬 공부한걸 사용해보고 싶어서, 리스트의 첫번째 값을 피벗으로 삼는 호어 분할 방식으로 시도해봤는데 계속 시간 초과가 난다. 퀵 정렬이 최악의 경우에는 시간 복잡도가 O(N²)라고 하더니 특정 케이스에서 제한 시간 내에 결과를 내지 못한 게 아닌가 싶다. 혹은 이보다 더 개선되도록 코드를 수정해야 한다거나... 

 

 

import sys
sys.setrecursionlimit(10**6)

def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    lesser_arr, equal_arr, greater_arr = [], [], []
    for num in arr:
        if num < pivot:
            lesser_arr.append(num)
        elif num > pivot:
            greater_arr.append(num)
        else:
            equal_arr.append(num)
    return quick_sort(lesser_arr) + equal_arr + quick_sort(greater_arr)

n, m = map(int, input().split())
temp = list(map(int, input().split()))
res = quick_sort(temp)
print(res[m-1])

 

찾아보니 리스트의 첫번째 데이터를 피벗으로 설정한 경우보다 중간 데이터를(index=len(arr)//2) 피벗으로 설정한 퀵 정렬 방식이 좀 더 빠르다고 하여 수정해본 코드. 이번엔 메모리 초과가 발생했다.

 

재귀 깊이 문제가 있나 싶어 코드를 추가하니까 시간초과가 발생한다.

일단 마무리하고... 조만간 다시 도전한다 퀵 정렬...