[백준] 2075번 n번째 큰 수, 파이썬

문제 바로가기

import sys
import heapq

n=int(input())


def heapsort(k):
  h=[]
  for i in range(n):
    arr=list(map(int,input().split()))
    if i==0:
      for j in arr:
        heapq.heappush(h,(-j,j))
    else:
      for j in arr:
        if j>h[n-1][1]:
          heapq.heappush(h,(-j,j))
        
  kth_max=None
  for _ in range(k):
    kth_max=heapq.heappop(h)[1]
  
    
  return kth_max

print(heapsort(n))

• python 3와 pypy3 모두 메모리 초과가 났던 코드이다. maxheap를 사용했는데 이 경우 힙의 길이가 너무 길어져 배열이 상당히 커지므로 메모리 초과가 난 것이다.

import sys
import heapq

n=int(input())
arr = [list(map(int, input().split())) for _ in range(n)]

def heapsort():
  h=[]
  for i in range(n):
    if i==0:
      for j in arr[i]:
        heapq.heappush(h,j)
    else:
      for j in arr[i]:
        if j>h[0]:
          heapq.heappush(h,j)
          heapq.heappop(h)
    
  return h[0]

print(heapsort())

• 힙의 크기를 일정하게 유지하기 위해 첫 배열 길이인 n에서 새로운 숫자가 추가될 때 마다 힙에서 가장 작은 수를 없애줬다. 최종 n 크기의 배열에서는 최종 상위 n개 수가 남게 되므로 가장 작은 수를 반환하면 답을 얻을 수 있다.
• python3 에서는 메모리 초과, pypy3에서는 정답이 났던 코드이다. 이 경우는 arr의 크기가 너무 커져서 메모리 초과가 났다.

import sys
import heapq

n=int(input())

def heapsort():
  h=[]
  for i in range(n):
    arr = list(map(int, input().split()))
    if not h:
      for i in arr:
        heapq.heappush(h,i)
    else:
      for i in arr:
        if i>h[0]:
          heapq.heappush(h,i)
          heapq.heappop(h)
    
  return h[0]

print(heapsort())

• 최종 정답 코드이다. arr을 반복문을 돌면서 다이렉트로 선언하니 arr의 배열 길이도 일정하고, 힙의 길이 또한 일정했다.