백준

[백준] 7662번: 이중 우선순위 큐(파이썬)

초코바나나쉐이크 2025. 1. 8. 18:12

https://www.acmicpc.net/problem/7662

 

우선순위 큐를 응용해서 푸는 문제였다.

heapq를 사용하여 구현하려고 했다.

 

import heapq

t = int(input())
for _ in range(t):
    k = int(input())

    dup = [] # 공용 리스트
    min_heap = [] # 최소힙
    max_heap = [] # 최대힙
    
    # k번 만큼 명령 실행
    for i in range(k):
        cmd, num = map(str, input().split())
        num = int(num)
        cnt = 0
        flag = True
        # 추가
        if cmd == "I":
            dup.append(num)
            heapq.heappush(min_heap, num)
            heapq.heappush(max_heap, -num)

        # 삭제
        elif cmd == "D":
            if dup:
            # 최소값 삭제
                if num == -1:
                    while flag:
                        if min_heap[0] not in dup:
                            heapq.heappop(min_heap)
                        else:
                            dup.remove(min_heap[0])
                            heapq.heappop(min_heap)
                            flag = False
                # 최대값 삭제
                if num == 1:
                    while flag:
                        if -max_heap[0] not in dup:
                            heapq.heappop(max_heap)
                        else:
                            dup.remove(-max_heap[0])
                            heapq.heappop(max_heap)
                            flag = False
    if dup:
        print(max(dup),min(dup))
    else:
        print("EMPTY")

최소힙, 최대힙, 공용리스트를 따로 만들어서 풀려고 했다.

D 1을 입력받아서 최대값을 삭제한다면, 최대힙에서만 삭제되고 최소힙에선 최댓값이 삭제된 것을 모른다.

그렇기 때문에 공용리스트를 만들어서 공용리스트에 그 값이 없다면 이미 지워진 것이므로 최소 힙에서도 최댓값을 지우도록 구현했다.

시간 초과..

입력값이 1,000,000개여서 시간 초과가 난다.

아마 min_heap[0 not in dup:과 -max_heap[0] not in dup: 이 부분에서 O(N)이어서 시간 초과가 나는 것 같다.

 

아무리 생각해도 O(logN)으로 줄일 방법이 생각나지 않아서 다른 분의 코드를 참고했다.

import heapq

def isEmpty(nums):
    # nums에 1인 데이터가 하나라도 있으면 비어있지 않은 것
    for item in nums:
        if item[1] > 0:
            return False
    return True

t = int(input())

for _ in range(t):
    k = int(input())

    nums = dict()
    min_heap = [] # 최소힙
    max_heap = [] # 최대힙
    
    # k번 만큼 명령 실행
    for i in range(k):
        cmd, num = map(str, input().split())
        num = int(num)
        
        # 추가
        if cmd == "I":
            # 중복 삽입일 때
            if num in nums:
                nums[num] += 1
            else:
                nums[num] = 1
                heapq.heappush(min_heap, num)
                heapq.heappush(max_heap, -num)

        # 삭제
        elif cmd == "D":
            if not isEmpty(nums.items()):
                # 최댓값 삭제
                if num == 1:
                    while -max_heap[0] not in nums or nums[-max_heap[0]] < 1:
                        temp = -heapq.heappop(max_heap)
                        if temp in nums:
                            del(nums[temp])
                    nums[-max_heap[0]] -= 1
                        
                # 최솟값 삭제
                else:
                    while min_heap[0] not in nums or nums[min_heap[0]] < 1:
                        temp = heapq.heappop(min_heap)
                        if temp in nums:
                            del(nums[temp])
                    nums[min_heap[0]] -= 1

    if isEmpty(nums.items()):
        print("EMPTY")
    else:
        while min_heap[0] not in nums or nums[min_heap[0]] < 1:
            heapq.heappop(min_heap)
        while -max_heap[0] not in nums or nums[-max_heap[0]] < 1:
            heapq.heappop(max_heap)
        print(-max_heap[0], min_heap[0])

리스트 대신 딕셔너리로 구현하였다.

 

드디어 Class 3++까지 모두 완료하였다.

아직 갈 길이 멀지만 계속 나아가자!!😎