[백준] 7662번: 이중 우선순위 큐

19377 단어 heapqpythonheapq

문제보기

나는 최소힙이랑 최대힙 두 개를 만들어서 풀려했다.

처음에 내가 한 방법은

  1. 입력이 들어오면 두 힙에 모두 push
  2. 한 쪽에서 빠지면 빠진 값 저장하기
  3. 나머지 한 쪽에서도 해당 값 지우기

였는데 3번이 찝찝했었는데 역시 시간초과가 났다.


그래서 다른 사람들이 한 방법을 찾아보고 내 코드를 바꿔봤다.

i 번째에 insert('I') 된 노드의 delete('D') 여부(boolean)를 저장하는 리스트를 모두 False로 초기화해놓고 나서

  1. push할 때 몇 번째 연산(i, 뒤에서는 idx로 쓰인다)인지와 함께 튜플로 힙에 저장했다.
	for i in range(t):
    	op, opn = input().split()
    	opn = int(opn)
    	if op == 'I':
        	heapq.heappush(minH, (opn, i))
        	heapq.heappush(maxH, (-opn, i))
  1. pop할 때 해당 노드가 이미 다른 힙에서 pop 되었었는지 확인(popped[idx]True 인지)
	# 최댓값 삭제
	if opn == 1:
    
    	# 이미 최소힙에서 pop된 애들일 경우 계속 삭제
    	while maxH and popped[maxH[0][1]]:
        	idx = heapq.heappop(maxH)[1]

    	if maxH:  # 힙이 비어있지 않으면
        	idx = heapq.heappop(maxH)[1]
        	popped[idx] = True

이렇게 하면 'D' 연산을 할 때 한 쪽 힙에서만 지워도 제대로 결과를 낼 수 있다! 밑에 코드를 보자

import heapq

T = int(input())
result = []
for _ in range(T):
    t = int(input())
    minH = []
    maxH = []
    popped = [False for _ in range(t)]
    for i in range(t):
        op, opn = input().split()
        opn = int(opn)
        if op == 'I':
            heapq.heappush(minH, (opn, i))
            heapq.heappush(maxH, (-opn, i))
        elif op == 'D':
            if opn == 1:
                while maxH and popped[maxH[0][1]]:
                    idx = heapq.heappop(maxH)[1]
                if maxH:
                    idx = heapq.heappop(maxH)[1]
                    popped[idx] = True
            else:
                while minH and popped[minH[0][1]]:
                    idx = heapq.heappop(minH)[1]
                if minH:
                    idx = heapq.heappop(minH)[1]
                    popped[idx] = True

    while minH and popped[minH[0][1]]:
        heapq.heappop(minH)
    while maxH and popped[maxH[0][1]]:
        heapq.heappop(maxH)
        
    if not minH:
        result.append('EMPTY')
    else:
        result.append(str(-heapq.heappop(maxH)[0])+' '+str(heapq.heappop(minH)[0]))

print('\n'.join(result))

주의해야 할 점은 연산을 모두 마치고 나서, 마지막 답으로 각 힙의 탑노드를 빼와야하기 때문에 마찬가지로 우선순위에 올라가있는 poppedTrue인 노드들을 먼저 삭제해주어야 한다.

대신 그냥 맨 위의 애들만 체크 하면 되기 때문에 따로 idx를 저장할 필요는 없고 삭제(heappop)만 해주면 된다.


하씨. 혹시나 해서 힙 한개만 써서 해봤는데 역시 시간초과다. 표준 입출력 써서 해두 시간초과다 파이썬 부들부들

# 시간초과
import heapq
import sys

input = sys.stdin.readline
print = sys.stdout.write

T = int(input())
result = []
for _ in range(T):
    t = int(input())
    h = []
    for i in range(t):
        op, opn = input().split()
        opn = int(opn)
        if op == 'I':
            heapq.heappush(h, opn)
        elif op == 'D':
            if h:
                if opn == 1:
                    h.pop()
                else:
                    h.pop(0)

    if not h:
        result.append('EMPTY')
    else:
        maxSol = h.pop()
        if h:
            minSol = h.pop(0)
        else:
            minSol = maxSol
        result.append(f'{maxSol} {minSol}')

print('\n'.join(result))
sys.stdout.flush()

좋은 웹페이지 즐겨찾기