[코드트리] [python] 삼성 2024 상반기 오전 2번 : 코드트리 투어

문제 설명

https://www.codetree.ai/ko/frequent-problems/problems/codetree-tour/submissions?page=1&page_size=20&introductionSetId=&bookmarkId=

삼성 코딩테스트 기출 문제 제출: 코드트리 투어 | 코드트리

문제 설명이 길어서 전체적인 문제 내용은 위의 링크를 따라 보면 된다. 이 글에서는 핵심 로직만 다루겠다.

 

 

문제 풀이 

1. 코드트리 랜드 건설

n, m = 0, 0
start = 0
graph = {}
products = {}
distances = []
product_heap = []

for c in command:
    if c[0] == 100:
        n, m = c[1], c[2]
        graph = {node: [] for node in range(n)}
        for i in range(3, len(c), 3):
            graph[c[i]].append((c[i + 1], c[i + 2]))
            graph[c[i + 1]].append((c[i], c[i + 2])) 
        dijkstra()
        rebuildHeap()  # 그래프 생성 후 한번만!

각 간선은 방향성을 갖지 않으므로, 출발 노드와 도착 노드 모두에 해당 간선 정보를 더해준다.

dictionary를 이용하여 구현하였다.

예시)

graph = {
    노드 번호: [ (연결된 노드 번호, 가중치), ... ],
    ...
}

graph = {
    0: [(1, 8), (2, 15), (3, 20), (4, 25), (5, 30), (1, 20)],
    1: [(0, 8), (0, 20), (2, 18), (4, 24), (5, 27), (4, 18), (4, 9), (3, 19)],
    2: [(0, 15), (1, 18), (3, 17)],
    3: [(0, 20), (1, 19), (2, 17), (4, 19), (5, 23)],
    4: [(0, 25), (1, 24), (1, 18), (1, 9), (3, 19), (5, 34)],
    5: [(0, 30), (5, 1), (1, 27), (3, 23), (4, 34)]
}

2. 다익스트라 알고리즘 with heap

먼저, 힙 자료구조 (우선순위 큐)를 사용하여, 최적 상품을 빠르게 찾았다.

import heapq

출발지에서 목적지까지 모든 노드까지의 최단 거리를 계산하기 위해, 위의 graph 정보를 이용하여 다익스트라 알고리즘을 사용한다.

다익스트라 알고리즘

하나의 노드에서 모든 다른 노드까지의 최단거리 구하기

• 최단거리를 구할 노드에서 시작하여, 거리가 입력된 노드 중 최단거리가 가장 작은 노드를 돌아가며 선택합니다.
• 노드를 돌아가면서, 더 거리가 짧은 나오면 값을 갱신하여 넣습니다.

우선순위 큐

• 언제든지 지금까지 발견한 가장 짧은 거리가 가장 먼저 처리되도록 관리
• 매 순간 "가장 유망한 노드"를 먼저 방문하기 위해

def dijkstra():
    global distances
    distances = [float('inf')] * n  # 출발지로부터 각 노드까지 거리 초기화
    distances[start] = 0  # 출발지까지 거리는 0
    queue = [(0, start)]  # 우선순위 큐: (거리, 노드)

    while queue:
        current_distance, current_node = heapq.heappop(queue)

        # 이미 더 짧은 거리로 방문했으면 스킵
        if current_distance > distances[current_node]:
            continue

        # 인접 노드 검사
        for neighbor, weight in graph[current_node]:
            distance = current_distance + weight
            if distance < distances[neighbor]:
                distances[neighbor] = distance
                heapq.heappush(queue, (distance, neighbor))

3. 여행상품 생성, 취소

    elif c[0] == 200:
        id_, revenue, dest = c[1], c[2], c[3]
        products[id_] = [revenue, dest]
        pushProduct(id_, revenue, dest)  # ✅ 바로 heap push (rebuild 안 함!)

    elif c[0] == 300:
        if c[1] in products:
            del products[c[1]]
        # lazy deletion

시간복잡도를 개선하기 위해 생성된 상품을 바로 heap에 넣는다.

product_heap = []

def pushProduct(id_, revenue, dest):
    cost = distances[dest]
    # 상품 배송 가능 여부 확인
    if cost != float('inf') and cost <= revenue:
        profit = revenue - cost
        heapq.heappush(product_heap, (-profit, id_, id_))

4. 최적의 여행 상품

위의 pushProduct() 함수를 보면,

생성된 상품의 revenue, dest, cost를 바로 계산하여 product_heap에 상품을 등록한다.

heapq는 최소 힙이므로 수익이 큰 상품이 먼저 나오게 하기 위해 수익성을 -profit으로 넣는다.

def getProduct():
    while product_heap:
        _, _, best_id = heapq.heappop(product_heap)
        if best_id in products:
            print(best_id)
            del products[best_id]
            return
    print(-1)

product_heap에서 가장 best인 여행상품을 출력하고,

없다면 (노드에 도착할 수 없다면) -1을 출력한다.

 

전체 코드

import heapq

def dijkstra():
    global distances
    distances = [float('inf')] * n
    distances[start] = 0
    queue = [(0, start)]

    while queue:
        current_distance, current_node = heapq.heappop(queue)

        if current_distance > distances[current_node]:
            continue

        for neighbor, weight in graph[current_node]:
            distance = current_distance + weight
            if distance < distances[neighbor]:
                distances[neighbor] = distance
                heapq.heappush(queue, (distance, neighbor))

def rebuildHeap():
    global product_heap
    product_heap = []
    for id_, (revenue, dest) in products.items():
        cost = distances[dest]
        if cost != float('inf') and cost <= revenue:
            profit = revenue - cost
            heapq.heappush(product_heap, (-profit, id_, id_))

def pushProduct(id_, revenue, dest):
    cost = distances[dest]
    if cost != float('inf') and cost <= revenue:
        profit = revenue - cost
        heapq.heappush(product_heap, (-profit, id_, id_))

def getProduct():
    while product_heap:
        _, _, best_id = heapq.heappop(product_heap)
        if best_id in products:
            print(best_id)
            del products[best_id]
            return
    print(-1)

# 메인
q = int(input())
command = []
for _ in range(q):
    c = list(map(int, input().split()))
    command.append(c)

n, m = 0, 0
start = 0
graph = {}
products = {}
distances = []
product_heap = []

for c in command:
    if c[0] == 100:
        n, m = c[1], c[2]
        graph = {node: [] for node in range(n)}
        for i in range(3, len(c), 3):
            graph[c[i]].append((c[i + 1], c[i + 2]))
            graph[c[i + 1]].append((c[i], c[i + 2])) 
        dijkstra()
        rebuildHeap()  # 그래프 생성 후 한번만!

    elif c[0] == 200:
        id_, revenue, dest = c[1], c[2], c[3]
        products[id_] = [revenue, dest]
        pushProduct(id_, revenue, dest)  # ✅ 바로 heap push (rebuild 안 함!)

    elif c[0] == 300:
        if c[1] in products:
            del products[c[1]]
        # lazy deletion

    elif c[0] == 400:
        getProduct()

    elif c[0] == 500:
        start = c[1]
        dijkstra()
        rebuildHeap()  # ✅ 출발지 바뀔 때만 rebuild

시간 초과 개선 전 코드

import heapq

def dijkstra():
    global distances
    distances = [float('inf')] * n
    distances[start] = 0
    queue = [(0, start)]

    while queue:
        current_distance, current_node = heapq.heappop(queue)

        if current_distance > distances[current_node]:
            continue

        for neighbor, weight in graph[current_node]:
            distance = current_distance + weight
            if distance < distances[neighbor]:
                distances[neighbor] = distance
                heapq.heappush(queue, (distance, neighbor))

def getProduct():
    candidates = []

    for id_, (revenue, dest) in products.items():
        cost = distances[dest]
        if cost != float('inf') and cost <= revenue:
            candidates.append((revenue - cost, id_, id_))

    if not candidates:
        print(-1)
        return

    candidates.sort(key=lambda x: (-x[0], x[1]))
    best_id = candidates[0][2]

    print(best_id)
    del products[best_id]

# 메인
q = int(input())
command = []
for _ in range(q):
    c = list(map(int, input().split(' ')))
    command.append(c)

n, m = 0, 0
start = 0  
graph = {}
products = {}
distances = []

for c in command:
    if c[0] == 100:
        n, m = c[1], c[2]
        graph = {node: [] for node in range(n)}
        for i in range(3, len(c), 3):
            graph[c[i]].append((c[i + 1], c[i + 2]))
            graph[c[i + 1]].append((c[i], c[i + 2])) 
        dijkstra()  # 그래프 생성 시

    elif c[0] == 200:
        products[c[1]] = [c[2], c[3]]

    elif c[0] == 300:
        if c[1] in products:
            del products[c[1]]
        # 상품 삭제 시 그래프는 안 바뀌므로 dijkstra 불필요

    elif c[0] == 400:
        getProduct()

    elif c[0] == 500:
        start = c[1]
        dijkstra()  # 출발지 변경 시

나는 그래프 생성 시, 출발지 변경 시 dijkstra()를 매번 실행하는 코드로 구현했었다.

gpt를 이용해서 heap 자료구조를 더욱 활용하여 개선해 보았다.