Archive / BOJ / Python

16455번 K번째 수 찾는 함수

이왜틀?

함수를 직접 구현하는 문제.

int n=a.length;
int start=0;
int end=n-1;
k-=1;

while(start<end) {
    int i=start;
    int j=end;
    int mid=a[(i+j)/2];
    
    while(i<j) {
        if(a[i]>=mid) {
            int tmp = a[j];
            a[j] = a[i];
            a[i] = tmp;
            j--;
        }else {
            i++;
        }
    }
    if(a[i]>mid) i--;
    if(k<=i) end = i;
    else start = i + 1;
}
return a[k];

퀵셀렉트를 했는데 시간초과가 났다.
조금만 수정하면 통과할 수 있다.

설계

java 코드에서 pivot을 중앙 값으로 설정했다.
그런데 찾는 값이 한쪽으로 치우치면 시간 복잡도가 O(N^2) 까지 갈 수 있다.

만약 pivot을 무작위 난수를 생성해 선택하면 한쪽으로 치우치는걸 다소 방지할 수 있다.
그리고 수열의 요소가 서로 다른 값임이 보장되지 않는다.
이 부분에 대한 로직도 추가한다.

구현

pivot은 start와 end 사이 랜덤으로 설정한다.
중복값인 경우 정확한 판단을 위해 pivot을 기준으로 3분할한다.

def kth(a, k):
    import random

    n = len(a)
    start = 0
    end = n - 1
    k -= 1

    while start <= end: 
        # 3분할
        i = start
        j = start
        l = end

        # 무작위 값
        pivot = a[random.randint(i, l)]

        while i <= l:
            # pivot보다 작을때
            if a[i] < pivot:
                a[i], a[j] = a[j], a[i]
                i += 1
                j += 1
            # pivot보다 클떄
            elif a[i] > pivot:
                a[i], a[l] = a[l], a[i]
                l -= 1
            else:
                i += 1
        
        
        if k < j:
            end = j - 1
        elif k > l:
            start = l + 1
        else:
            return pivot

채점

반성

좀 빡쎘다…
옛날 코드 잘 해놔서 그대로 포팅하면 쉽게 날먹할 줄 알았는데

코드 확인

Link to GitHub