24173 힙 정렬 - 시간줄이기

문제

오늘도 서준이는 최소 힙 기반 힙 정렬 수업 조교를 하고 있다. 아빠가 수업한 내용을 학생들이 잘 이해했는지 문제를 통해서 확인해보자.

N개의 서로 다른 양의 정수가 저장된 배열 A가 있다. 힙 정렬로 배열 A를 정렬할 경우 배열 A에 K 번째 교환되는 수를 구해서 우리 서준이를 도와주자.

크기가 N인 배열에 대한 힙 정렬 의사 코드는 다음과 같다.

heap_sort(A[1..n]) { # A[1..n]을 정렬한다.
    build_min_heap(A, n);
    for i <- n downto 2 {
        A[1] <-> A[i];  # 원소 교환
        heapify(A, 1, i - 1);
    }
}

build_min_heap(A[], n) {
    for i <- ⌊n / 2⌋ downto 1
        heapify(A, i, n)
}

# A[k]를 루트로 하는 트리를 최소 힙 성질을 만족하도록 수정한다.
# A[k]의 두 자식을 루트로 하는 서브 트리는 최소 힙 성질을 만족하고 있다.
# n은 배열 A의 전체 크기이며 최대 인덱스를 나타낸다.
heapify(A[], k, n) {
    left <- 2k; right <- 2k + 1;
    if (right ≤ n) then {
        if (A[left] < A[right]) then smaller <- left;
                                else smaller <- right;
    }
    else if (left ≤ n) then smaller <- left;
    else return;

    # 최소 힙 성질을 만족하지 못하는 경우 재귀적으로 수정한다.
    if (A[smaller] < A[k]) then {
        A[k] <-> A[smaller];
        heapify(A, smaller, n);
    }
}

입력

첫째 줄에 배열 A의 크기 N(5 ≤ N ≤ 500,000), 교환 횟수 K(1 ≤ K ≤ 108)가 주어진다.

다음 줄에 서로 다른 배열 A의 원소 A1, A2, ..., AN이 주어진다. (1 ≤ Ai ≤ 109)

5 2
2 5 1 4 3

출력

K 번째 교환되는 두 개의 수를 작은 수부터 한 줄에 출력한다. 교환 횟수가 K 보다 작으면 -1을 출력한다.

1 2

문제풀이

시간 줄이는 게 가장 중요한 문제였다!!

힙 정렬을 하되, K번째 교환하는 값을 출력하는 것이 문제의 최종 목적이었다.

나는 너무 정직하게 힙 정렬을 모두하고, 교환 값을 리스트에 갱신했다.

그런데  K개의 크기: (1 ≤ K ≤ 10^8) 배열의 크기 : (1 ≤ Ai ≤ 10^9)로 모두 정렬을 하고, 리스트에 모두 갱신하면 시간 초과가 났다.

• 교환 값을 리스트에 담는 것은 시간이 너무너무 많이 걸린다!! 이전 값이 필요한게 아니라면 최대한 리스트를 쓰지 말자.
• K번째 값이 나왔다면 힙을 더 정렬할 필요가 없다!! K 번째가 될 때 분기점을 찾아보자..!!

# 24713 힙정렬
import sys

COUNT = 0
exchange_list = []

def heapify(a_list,root,n):
    global K
    global COUNT
    #부모노드가 최소가 되도록
    child = root*2 + 1
    while child < n:
        if child + 1 < n:
            if a_list[child] > a_list[child+1]:
                child += 1
        if a_list[root] > a_list[child]:
            COUNT += 1
            if COUNT == K:
                exchange_list.append([a_list[child],a_list[root]])
                return
            a_list[root],a_list[child] = a_list[child],a_list[root]
            root = child
            child = root * 2 + 1
        else:
            break
    return a_list

def heap_sort(a_list):
    global K
    global COUNT
    n = len(a_list)
    root = int((n-1)/2)
    global exchange_list
    while root >= 0:
        a_list = heapify(a_list,root,n)
        if COUNT == K:
            return
        root += -1

    i = n-1

    while i >= 1 :
        COUNT += 1
        if COUNT == K:
            exchange_list.append([a_list[0], a_list[i]])
            return
        a_list[0],a_list[i] = a_list[i],a_list[0]
        a_list = heapify(a_list,0,i)
        if COUNT == K:
            return
        i += -1

N, K = map(int, input().split())
a_list =  list(map(int, input().split()))

heap_sort(a_list)
#
# print(a_list)
# print(exchange_list)

if K > COUNT:
    print(-1)
else:
    print(' '.join(list(map(str,exchange_list[0]))))