[Python] 코드트리 투어: 삼성SW역량테스트 2024 상반기 오전 2번

 

코드트리 투어: 삼성SW역량테스트 2024 상반기 오전 2번

 

Information

• 자료구조: 그래프, 힙(우선순위 큐)
• 알고리즘: 다익스트라

Query

이 문제는 다양한 쿼리가 인풋으로 들어온다.

그래프를 만드는 쿼리 빼고 4개의 쿼리에 대해 처리해야한다.

• (1) 랜드 건설: 가중치 그래프를 만드는 쿼리이다. 맨 처음 인풋으로만 들어온다.
• (2) 여행 상품 생성: 현재 출발지 기준으로 매출, 도착지 정보를 바탕으로 수익을 계산한다.
• (3) 여행 상품 취소: 특정 상품 id가 (2)로 생성되어있다면, 취소한다.
• (4) 최적 상품 판매: 수익이 가장 큰 상품을 판매하고, 그러한 상품이 없다면 -1을 출력
• (5) 상품 출발지 변경: 출발지를 새로운 출발지로 변경한다. (2)에서 생성된 정보 업데이트가 필요하다.

Setup

필요한 전역변수와 함수의 프로토타입을 정의하자.

import heapq
from collections import defaultdict

INF = float('inf')
n, m = 0, 0
adj = []
cost = []
item_dict = defaultdict(int)

 

`INF`는 다익스트라 알고리즘에서 사용될 상수이다.

`n`과 `m`은 프로그램 인풋으로 한 번 주어지고, 각각 노드 개수, 엣지 개수이다.

`adj`는 인접리스트이다

`cost`는 다익스트라 알고리즘에서 사용될 S(출발지)로부터 거리이다.

`item_dict`는 여행상품이 생성되면 저장할 딕셔너리이다.

 

def dijkstra(start):
    global n, cost, adj
    pass
    
def sell(pq):
    global item_dict
    pass
    
def update_item(pq):
    global cost
    pass

 

`dijkstra(start)`는 `start`를 출발지로 하는 다익스트라 알고리즘이다. 

맨 처음, 그리고 (5)번 쿼리에서 사용할 것이다.

우선순위큐를 이용한 알고리즘의 시간복잡도는 $O(E \log V)$이지만, 문제에서 최대 15번 쿼리가 들어오므로 괜찮을 것으로 보인다. 

 

`sell(pq)`는 (4)번 쿼리에 대한 함수이다.

 

`update_item(pq)`는 (5)번 쿼리에 대한 함수이다. 여기서의 `pq`는 여행상품에 대한 우선순위큐이다.

 

 

이제 여행상품을 정의해보자.

정렬될 상태가 필요하고, 정렬 순서는 (수익이 클 수록, id가 작을 수록)이다.

참고로 (수익) = (매출) - (비용)이며, 도착지에 도달할 수 없거나 수익이 음수인 경우 판매 불가능하다.

그리고 (5)번 쿼리마다 cost[]가 바뀌고 이에 따라 (수익)도 달라지므로 재정렬할 필요가 있다.

`cost` 배열이 바뀌려면 도착지(`dest`) 정보도 버리지 않아야겠다.

여행상품은 `revenue, dest, id, profit`을 포함해야겠다.

튜플로 저장해서 표현할 수 있지만, class 객체로 구현해보자.

class Item:
    def __init__(self, item_id, revenue, dest):
        self.item_id = item_id
        self.revenue = revenue
        self.dest = dest
        self.profit = 0

        self.update_profit()

    def __lt__(self, other):
        if self.profit == other.profit:
            return self.item_id < other.item_id
        return self.profit > other.profit

    def update_profit(self):
        # revenue - cost
        global cost
        if cost[self.dest] == INF:
            self.profit = -1
        else:
            self.profit = self.revenue - cost[self.dest]

 

`__lt__(self, other)`로 `Item` 객체의 정렬 순서를 재정의할 수 있다.

 

`update_profit()`로 (5)번 쿼리에 대한 대응코드를 작성했다.

어차피 `profit`이 음수이면 판매 가능하므로, 도달할 수 없는 경우에도 임의의 음수 (`-1`)로 할당했다.

(1) 랜드 건설

def main():
    global Q, n, m, adj, cost, item_dict
    Q = int(input())
    cmds = list(map(int, input().split()))
    n, m = cmds[1], cmds[2]
    adj = [[] for _ in range(n)]
    for i in range(3, len(cmds), 3):
        v, u, w = cmds[i:i + 3]
        adj[v].append((u, w))
        adj[u].append((v, w))

    dijkstra(0)
    pq_item = []

 

쿼리를 한줄로 받고, 모든 값이 정수임이 보장되므로 리스트로 저장한다.

`n`과 `m`을 따로 추출하고 `adj`를 만들자

3칸씩 인덱스를 증가시켜서 가중치 그래프를 표현하는 인접리스트를 만든다.

그래프가 완성되었다면 출발지를 `0`으로 하는 `cost` 배열을 채운다.

 

`pq_item`은 여행상품에 대한 우선순위큐이다.

 

(2) 여행 상품 생성

200으로 시작하는 쿼리이다.

def main():
    global Q, n, m, adj, cost, item_dict
    #(...)
    for _ in range(Q - 1):
        cmd = list(map(int, input().split()))
        if cmd[0] == 200:
            _, item_id, revenue, dest = cmd
            item = Item(item_id, revenue, dest)
            heapq.heappush(pq_item, item)
            item_dict[item_id] = 1

 

상품id, 매출, 목적지를 입력받고, 해당 정보를 포함하는 `Item` 객체를 생성한다.

`pq_item`에 push한다.

그리고 `item_dict`에 생성되었다는 표시를 `1`로 한다.

(3) 여행 상품 취소

단순히 `item_dict`에서 지워버린다.

`defaultdict`를 이용하여 0으로 바꾸었다.

내장 객체 `dict`를 이용하여 `del item_dict[item_id]`로 처리할 수도 있다.

def main():
    global Q, n, m, adj, cost, item_dict
    #(...)
    for _ in range(Q - 1):
        cmd = list(map(int, input().split()))
        if cmd[0] == 200:
            pass
        elif cmd[0] == 300:
            _, item_id = cmd
            item_dict[item_id] = 0

 

(4) 최적 상품 판매

이 프로그램의 유일한 출력이 필요한 쿼리이다.

`sell()` 함수를 통해 정답을 출력하자.

def main():
    global Q, n, m, adj, cost, item_dict
    #(...)
    for _ in range(Q - 1):
        cmd = list(map(int, input().split()))
        if cmd[0] == 200:
            pass
        elif cmd[0] == 300:
            pass
        elif cmd[0] == 400:
            ret = sell(pq_item)
            print(ret)

 

def sell(pq):
    global item_dict

    while pq:
        if pq[0].profit < 0:
            break
        item = heapq.heappop(pq)
        if item_dict[item.item_id] == 1:
            return item.item_id
    return -1

 

상품 우선순위 큐에서 차례대로 첫번째 원소에 접근한다.

첫번째 상품의 `profit`이 음수라면 판매 불가능하므로 반복문을 `break`한다

그게 아니라면 `pop`하여 해당 상품의 `id` 정보를 가져온다.

상품 `id`가 취소되지 않았다면 `item_dict`에서 1을 value로 갖는다. 

이 경우에는 `item_id`를 return하고, `item_dict`에서 `1`의 값이 아니라면 `-1`을 `return`한다.

(5) 출발지 변경

`dijkstra(), update_item()` 두번 호출하여 여행상품의 정보를 업데이트한다.

def main():
    global Q, n, m, adj, cost, item_dict
    #(...)
    for _ in range(Q - 1):
        cmd = list(map(int, input().split()))
        if cmd[0] == 200:
            pass
        elif cmd[0] == 300:
            pass
        elif cmd[0] == 400:
            pass
        elif cmd[0] == 500:
            _, s = cmd
            dijkstra(s)
            update_item(pq_item)

 

def dijkstra(start):
    global n, cost, adj
    cost = [INF] * n
    cost[start] = 0
    pq = [(0, start)]

    while pq:
        cur_cost, cur = heapq.heappop(pq)
        if cur_cost > cost[cur]:
            continue
        for nxt, w in adj[cur]:
            nxt_cost = cost[cur] + w
            if nxt_cost < cost[nxt]:
                cost[nxt] = nxt_cost
                heapq.heappush(pq, (nxt_cost, nxt))

 

다익스트라 알고리즘은 우선순위큐를 이용하는 알고리즘을 이용했다.

일반적인 코드이다.

def update_item(pq):
    global cost
    for i in range(len(pq)):
        pq[i].update_profit()
    heapq.heapify(pq)

 

`pq_tmp`를 만들수도있지만, 그냥 기존 `pq_item`의 인덱스로 직접 접근하였다.

마지막 `heapify()`는 $O(N)$이므로 선형시간안에 업데이트할 수 있다.

참고로 비어있는 `pq`에 $N$개의 원소를 `push`하면 $O(N \log N)$이다.

 

Python Code

주석없는 코드로 100줄 나왔다.

삼성SW역량테스트 문제 치고 코드가 짧은 편인 것 같다.

import heapq
from collections import defaultdict

INF = float('inf')
n, m = 0, 0
adj = []
cost = []
item_dict = defaultdict(int)

class Item:
    def __init__(self, item_id, revenue, dest):
        self.item_id = item_id
        self.revenue = revenue
        self.dest = dest
        self.profit = 0

        self.update_profit()

    def __lt__(self, other):
        if self.profit == other.profit:
            return self.item_id < other.item_id
        return self.profit > other.profit

    def update_profit(self):
        # revenue - cost
        global cost
        if cost[self.dest] == INF:
            self.profit = -1
        else:
            self.profit = self.revenue - cost[self.dest]

def dijkstra(start):
    global n, cost, adj
    cost = [INF] * n
    cost[start] = 0
    pq = [(0, start)]

    while pq:
        cur_cost, cur = heapq.heappop(pq)
        if cur_cost > cost[cur]:
            continue
        for nxt, w in adj[cur]:
            nxt_cost = cost[cur] + w
            if nxt_cost < cost[nxt]:
                cost[nxt] = nxt_cost
                heapq.heappush(pq, (nxt_cost, nxt))

def sell(pq):
    global item_dict

    while pq:
        if pq[0].profit < 0:
            break
        item = heapq.heappop(pq)
        if item_dict[item.item_id] == 1:
            return item.item_id
    return -1

def update_item(pq):
    global cost
    for i in range(len(pq)):
        pq[i].update_profit()
    heapq.heapify(pq)

def main():
    global Q, n, m, adj, cost, item_dict
    Q = int(input())
    cmds = list(map(int, input().split()))
    n, m = cmds[1], cmds[2]
    adj = [[] for _ in range(n)]
    for i in range(3, len(cmds), 3):
        v, u, w = cmds[i:i + 3]
        adj[v].append((u, w))
        adj[u].append((v, w))

    dijkstra(0)
    pq_item = []

    for _ in range(Q - 1):
        cmd = list(map(int, input().split()))
        if cmd[0] == 200:
            _, item_id, revenue, dest = cmd
            item = Item(item_id, revenue, dest)
            heapq.heappush(pq_item, item)
            item_dict[item_id] = 1
                
        elif cmd[0] == 300:
            _, item_id = cmd
            item_dict[item_id] = 0
        elif cmd[0] == 400:
            ret = sell(pq_item)
            print(ret)
        elif cmd[0] == 500:
            _, s = cmd
            dijkstra(s)
            update_item(pq_item)
        #print(item_dict)

if __name__ == '__main__':
    main()