N번째 큰 수

문제

내가 생각했던 답(메모리 초과)

예제를 기준으로 설명하면
arr[1] = [12, 13, 21, 48, 52]
arr[2] = [7, 8, 10, 14, 20]
arr[3] = [9, 11, 26, 28, 32]
arr[4] = [15, 19, 31, 35, 41]
arr[5] = [5, 6, 16, 25, 49]
로 입력받을 수 있다.

항상 arr[i]의 마지막 값은 그 배열의 최대값인 것을 알 수 있다.
그래서 찾고자 하는 N번째 큰 수를 구하기 위해서
N만큼 배열의 최대값을 비교를 한다.

즉 52, 20, 32, 41, 49 중 가장 큰 값은 arr[1]의 52이다.
그러면 arr[1].pop을 해주면 arr[1] = [12, 13, 21, 48]이 왼다.

이제 두 번째로 큰 값을 구하면 48, 20, 32, 41, 49이다.
여기서 가장 큰 값은 arr[5] = 49이므로 arr[5].pop을 해주면 arr[5] = [5, 6, 16, 25]가 된다. 이런식으로 N번째 까지 찾는 로직이였다.

import sys

for test_case in range(1):
    n = int(sys.stdin.readline())

    arr = [[] for _ in range(n)]
    for _ in range(n):
        tmp = list(map(int, sys.stdin.readline().split()))
        for i in range(n):
            arr[i].append(tmp[i])

    for _ in range(n):
        cur_max = -float("inf")
        cur_idx = 0
        for i in range(n):
            tmp = arr[i]
            if tmp:
                if arr[i][-1] > cur_max:
                    cur_max = arr[i][-1]
                    cur_idx = i

        arr[cur_idx].pop()
    print(cur_max)

하지만 -10억 ~ 10억이 이라는 큰 값이 배열에 저장되어 있어서 메모리 초과가 발생한다.

최적화

힙을 사용한다.
N번째 최대값을 구하는 것이므로 최대힙이 아닌 최소힙을 사용한다.
그리고 N이 넘어가면 heappop을 해준다.
왜냐하면 최대힙을 사용하면 [52, ...]가 되지만
최소힙을 사용하면 [35, ....]가 된다.

for test_case in range(1):
    n = int(sys.stdin.readline())

    arr = []

    for _ in range(n):
        tmp = list(map(int, sys.stdin.readline().split()))
        for i in range(n):
            heappush(arr, tmp[i])
            if len(arr) > n:
                heappop(arr)

    print(arr[0])