[Python] 코드트리 / 코드트리 투어

문제 링크 <- 진짜 좋은 문제...

 

객체를 heapq로 관리하려면...

진짜 도저히 안 풀려서 정답 코드를 참고하였다. 즉, 밑에 작성할 풀이 코드는 정답 코드에서 상당 부분을 참고했음을 밝힌다.

 

실전 코테에서 class로 객체를 선언해서 문제를 푸는 것은, 사실 실현하기는 어려울 것이다. 어지간히 능숙하지 않고서야 더 헷갈리기 마련이니까. 하지만, 조금이라도 시간복잡도 갖고 장난치는 문제를 만나면 객체를 선언하는 방법 말고는 없을 때가 있다.

 

그런데 문제가 있었다. 나는 지금까지 객체를 heapq로 관리할 줄 몰랐어서 heapq를 사용해야 할 때면 객체 사용을 포기하고는 했다. 그런데 정답 코드를 공부하면서 객체를 heapq로 관리할 수 있는 방법을 터득하였다!

class Package:
    def __init__(self, id_, revenue, dest, profit):
        self.id_ = id_
        self.revenue = revenue
        self.dest = dest
        self.profit = profit
    
    def __lt__(self, other):  # min heap에서 self가 other보다 우선시되려면:
        if self.profit == other.profit:
            return self.id_ < other.id_ 
        return self.profit > other.profit

 

여행 패키지를 관리하는 문제이므로, 당연히 패키지 객체를 만드는 것이 좋다. 그런데 패키지를 리스트로 묶어서 heapq로 관리하고자 할 때, 만약 __lt__를 해놓지 않는다면 에러가 뜰 것이다. 크기를 비교할 수 없기 때문이다. 하지만 __lt__를 통해 크기 비교의 기준을 우리가 정해준다면, heapq를 문제 없이 사용할 수 있다.

 

더 좋은 것은, __lt__의 결과를 고의적으로 반대로 지정함으로써 MAX Heap을 간접적으로 구현할 수 있다는 것이다!

 

외에도,

•  __lt__(self, other) : < 
•  __le__(self, other) : <= 
•  __gt__(self, other) : >
•  __ge__(self, other) : >= 
•  __eq__(self, other) : == 
•  __ne__(self, other) : != 
•  __add__(self, other) : + 
•  __sub__(self, other) : - 
•  __mul__(self, other) : *
•  __truediv__(self, other) : /
•  __floordiv__(self, other) : // 
•  __mod__(self, other) : % 
•  __pow__(self, other) : **

등이 있다.

 

정답 코드

import heapq
from collections import defaultdict
import sys
input = sys.stdin.readline

INF = float('inf')

class Package:
    def __init__(self, id_, revenue, dest, profit):
        self.id_ = id_
        self.revenue = revenue
        self.dest = dest
        self.profit = profit
    
    def __lt__(self, other):  # min heap에서 self가 other보다 우선시되려면:
        if self.profit == other.profit:
            return self.id_ < other.id_ 
        return self.profit > other.profit 

def dijkstra(s, graph):
    n = len(graph) 
    dist = [INF] * n
    dist[s] = 0
    visited = [False] * n

    for _ in range(n):
        node = -1
        min_dist = INF
        for j in range(n):
            if not visited[j] and min_dist > dist[j]:
                min_dist = dist[j]
                node = j
        if node == -1:
            break
        visited[node] = True
        for j in range(n):
            if graph[node][j] != INF and dist[j] > dist[node] + graph[node][j]:
                dist[j] = dist[node] + graph[node][j]

    return dist

Q = int(input())

queue = []
id_exist = defaultdict(bool)
dep = 0
started = True

for _ in range(Q):
    order = list(map(int, input().split()))

    if order[0] == 100:
        n = order[1]
        m = order[2]
        graph = [[INF] * n for _ in range(n)]
        for i in range(n):
            graph[i][i] = 0
        for i in range(1, m + 1):
            u, v, w = order[3 * i], order[3 * i + 1], order[3 * i + 2]
            val = min(w, graph[u][v], graph[v][u])
            graph[u][v] = val
            graph[v][u] = val
    
    elif order[0] == 200:
        if started:
            dist = dijkstra(dep, graph) 
            started = False
        _, id_, revenue, dest = order
        if not id_exist[id_]:  # 중복된 ID 방지
            id_exist[id_] = True
            profit = revenue - dist[dest]
            package = Package(id_, revenue, dest, profit)
            heapq.heappush(queue, package)

    elif order[0] == 300:
        id_ = order[1]
        id_exist[id_] = False  # 해당 ID의 패키지를 비활성화
        # queue에는 패키지를 남겨두지만, 출력 시 체크하여 제외
        
    elif order[0] == 400:
        #print([(package.id_, package.profit) for package in queue])
        while queue:
            priority = queue[0]
            if priority.profit < 0:
                print(-1)
                break
            heapq.heappop(queue)
            id_ = priority.id_
            if id_exist[id_]:
                print(id_)
                break
        else: 
            print(-1)
    
    elif order[0] == 500:
        dep = order[1]
        dist = dijkstra(dep, graph) 

        for package in queue:
            if id_exist[package.id_]:  # 활성화된 패키지만 업데이트
                package.profit = package.revenue - dist[package.dest]
        heapq.heapify(queue)  # 다시 힙 정렬해야함