[선택] 백준 11004번 K번째 수 시간초과 에러

신찬수 교수님의 알고리즘 유투브 영상으로 quick selection 알고리즘과 median of median 알고리즘을 배웠습니다.

배운 내용을 적용하고 싶어서 백준에서 quick selection 알고리즘으로 풀 수 있는 K 번째 수 문제를 도전해보았습니다!

 

11004번: K번째 수

예제 테스트 케이스의 경우에는 정상적으로 출력되고,  quick selection 알고리즘을 구현하여 문제를 풀었는데 시간 초과 에러가 발생했어요ㅠㅠ

시간 초과 에러ㅠㅠ

혼자 끙끙거리다가 구글에 이 문제를 푼 사람들의 코드를 보았는데, 대부분의 사람들이 sort 함수를 사용해서 풀었는데요,,,

저는 sort 함수를 그대로 사용해서 풀고 싶지는 않아서 무엇이 문제인지 고민해야겠습니다..

교수님께서 quick selection 알고리즘은 피봇을 잘 못 선택하여 분류한 두 그룹에 모든 수가 몰리는 worst case의 경우, O(n^2) 의 시간이 소요된다고 하셨습니다.

아무래도 백준이 평가하는 테스트케이스에 이런 worst case 가 포함되어 있어 시간 초과 에러가 발생한 것으로 생각됩니다,,,

내일 한 번 지금 코드를 어떻게 개선하면 좋을지 생각해보고, median of median 알고리즘을 통해서도 풀어봐야겠어요!

 

 

아래는 시간 초과 에러가 난 코드입니다ㅠㅠ

#백준 11004번 K번째 수

def quickselect(L, k):
    # pivot 설정하기
    p = L[0]
    # A: pivot 보다 작은 수들의 리스트
    A =[]
    # B: pivot 보다 큰 수들의 리스트
    B = []
    # M: pivot 과 같은 수들의 리스트
    M = []
    for i in L:
        if i > p:
            B.append(i)
        elif i < p:
            A.append(i)
        else:
            M.append(i)
    #찾고자 하는 k번째 수가 A의 길이보다 작은 경우: k번째 수는 A 안에 있음        
    if len(A) >= k:
        return quickselect(A, k)
    
    #찾고자 하는 k번째 수가 A와 M의 길이의 합보다 큰 경우: k번째 수는 B 안에 있음        
    elif len(A) + len(M) < k:
        return quickselect(B, k-len(A)-len(M))
    
    #위의 경우에 모두 해당이 안 되는 경우 pivot이 k번째 수임
    else:
        return p

total, kth = map(int,input("").split())

element = map(int,input("").split())
L = list(element)
print(quickselect(L, kth))

 

+2020.08.05

 

테스트케이스에 입력값이 굉장히 많은 것으로 예상되어 보니, input 함수로 입력을 받게 되면 그 자체에서 시간이 오래 걸린다고 합니다. 

아래 링크에 보면 sys.stdin.readline()을 사용하는 것이 시간 단축에 효율적이라고 쓰여있어, 수정하였습니다.

https://www.acmicpc.net/problem/15552

그리고 아무래도 pivot을 잘 못 설정한 것이 문제인 것 같아 pivot을 효율적으로 선택할 수 있는 median of three 알고리즘을 구현하려고 했습니다.

랜덤으로 하기 전에 그냥 앞에 세 개의 수로 설정하였더니, 똑같이 시간 초과 에러가 났습니다.

랜덤으로 구현해도,, 같은 에러가 나네요ㅠㅠ

제가 통과하지 못한 테스트케이스에는 역순으로 있는 것 같네요..ㅎㅎ 

아래는 여전히 시간초과 에러가 나는 수정한 코드입니다ㅠㅠ

#백준 11004번 K번째 수
import sys

def quickselect(L, k ):
    p = L[0]
    # A: pivot 보다 작은 수들의 리스트
    A =[]
    # B: pivot 보다 큰 수들의 리스트
    B = []
    # M: pivot 과 같은 수들의 리스트
    M = []
    for i in L:
        if i > p:
            B.append(i)
        elif i < p:
            A.append(i)
        else:
            M.append(i)
    #찾고자 하는 k번 째 수가 A의 길이보다 작은 경우: k번 째 수는 A 안에 있음
    if len(A) >= k:
        return quickselect(A, k )

    #찾고자 하는 k번 째 수가 A와 M의 길이의 합보다 큰 경우: k번 째 수는 B 안에 있음
    elif len(A) + len(M) < k:
        return quickselect(B, k-len(A)-len(M))

    #위의 경우에 모두 해당이 안 되는 경우 pivot이 k번 째 수임
    else:
        return p
def swap (L, a, b):
    tmp = L[a]
    L[a] = L[b]
    L[b] = tmp
    return L

def revised_list(L):
    if L[0]> L[1]:
        swap(L, 0, 1)
    elif L[1]>L[2]:
        swap(L, 1, 2)
    elif L[0]>L[2]:
        swap(L, 0, 2)
    return swap(L, 0, 1)

#input 대신 sys 모듈 사용
total, kth = map(int, sys.stdin.readline().split())

L = list(map(int, sys.stdin.readline().split()))
if total > 5:
    L = revised_list(L)

print(quickselect(L, kth))