연습 문제

바로 적용하는 문제부터 코딩 테스트에 자주 나오는 모델링 기법까지 네 문제를 풀어 봅니다. 접근 방법을 읽기 전에 먼저 직접 풀어 보세요. 모든 풀이는 Python 표준 라이브러리만 씁니다.

문제 1: 전파 시간

서버 n대가 0..n-1로 번호가 매겨져 있고 단방향 링크 (u, v, t)가 있습니다. u에서 보낸 메시지는 t밀리초 뒤 v에 닿습니다. 서버 s가 방송을 시작하고, 각 서버는 메시지를 받는 즉시 전달합니다. 마지막 서버가 메시지를 받는 시각을 구하고, 끝내 받지 못하는 서버가 있으면 -1을 돌려줍니다.

접근: 각 서버의 도착 시각은 s로부터의 최단 거리입니다. 다익스트라를 한 번 돌리고 최댓값을 구하면 되며, 무한대가 남아 있으면 닿지 않는 서버가 있는 것입니다.

import heapq
from math import inf

def broadcast_time(n, links, s):
    adj = [[] for _ in range(n)]
    for u, v, t in links:
        adj[u].append((v, t))
    dist = [inf] * n
    dist[s] = 0
    heap = [(0, s)]
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist[u]:
            continue
        for v, t in adj[u]:
            if d + t < dist[v]:
                dist[v] = d + t
                heapq.heappush(heap, (d + t, v))
    worst = max(dist)
    return -1 if worst == inf else worst

print(broadcast_time(4, [(0, 1, 2), (0, 2, 5), (1, 2, 1), (2, 3, 2)], 0))  # 5

문제 2: 할인권 한 장

도시 0에서 도시 n - 1까지 양방향 도로 (u, v, c)로 이동합니다. 도로 하나의 요금을 절반(내림)으로 만드는 할인권이 한 장 있습니다. 최소 총비용은 얼마일까요?

접근: 상태를 (도시, 할인권 사용 여부)로 넓힙니다. 상태 그래프에는 정점이 2n개 있고, 도로마다 사용 여부를 그대로 두는 간선이 생기며, 아직 할인권을 쓰지 않았다면 반값으로 가면서 사용 여부를 켜는 간선도 생깁니다. 이 상태 그래프에서 다익스트라를 돌리면 답이 나옵니다. "최대 k개 간선 무료" 같은 변형도 k + 1개 층으로 같은 방식으로 풉니다.

import heapq
from math import inf

def cheapest_trip(n, roads):
    adj = [[] for _ in range(n)]
    for u, v, c in roads:
        adj[u].append((v, c))
        adj[v].append((u, c))
    dist = [[inf, inf] for _ in range(n)]
    dist[0][0] = 0
    heap = [(0, 0, 0)]  # (비용, 도시, 사용 여부)
    while heap:
        d, u, used = heapq.heappop(heap)
        if d > dist[u][used]:
            continue
        for v, c in adj[u]:
            moves = [(d + c, used)]
            if not used:
                moves.append((d + c // 2, 1))
            for nd, nu in moves:
                if nd < dist[v][nu]:
                    dist[v][nu] = nd
                    heapq.heappush(heap, (nd, v, nu))
    return min(dist[n - 1])

print(cheapest_trip(4, [(0, 1, 10), (1, 3, 10), (0, 2, 3), (2, 3, 30)]))  # 15

문제 3: 가장 적게 벽 부수기

격자에 .(빈칸)과 #(벽)이 있습니다. 빈칸인 왼쪽 위 칸에서 출발해 네 방향으로 움직이며, 벽 칸에 들어가려면 그 벽을 부숴야 합니다. 오른쪽 아래 칸에 닿으려면 벽을 최소 몇 개 부숴야 할까요?

접근: 빈칸에 들어가는 비용은 0, 벽에 들어가는 비용은 1이므로 0-1 최단 경로 문제입니다. 힙 대신 덱을 씁니다.

from collections import deque
from math import inf

def fewest_walls(grid):
    h, w = len(grid), len(grid[0])
    dist = [[inf] * w for _ in range(h)]
    dist[0][0] = 0
    dq = deque([(0, 0)])
    while dq:
        r, c = dq.popleft()
        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nr, nc = r + dr, c + dc
            if 0 <= nr < h and 0 <= nc < w:
                cost = 1 if grid[nr][nc] == "#" else 0
                if dist[r][c] + cost < dist[nr][nc]:
                    dist[nr][nc] = dist[r][c] + cost
                    if cost == 0:
                        dq.appendleft((nr, nc))
                    else:
                        dq.append((nr, nc))
    return dist[h - 1][w - 1]

print(fewest_walls(["..#.", "###.", "...#", ".##."]))  # 2

문제 4: 우물 잇는 배관

우물들이 정수 좌표에 있습니다. 두 우물을 잇는 배관의 비용은 두 점의 맨해튼 거리입니다. 모든 우물을 (직접 또는 다른 우물을 거쳐) 최소 총비용으로 연결하세요.

접근: 모든 쌍이 간선 후보이므로 간선이 약 n^2 / 2개인 완전 그래프입니다. 크루스칼은 이 간선을 모두 정렬해야 하지만, 배열을 쓰는 프림은 간선 리스트를 만들지 않고 O(n^2)에 끝납니다.

from math import inf

def pipe_cost(points):
    n = len(points)
    best = [inf] * n  # 지금까지의 트리와 각 우물을 잇는 가장 싼 배관
    in_tree = [False] * n
    best[0] = 0
    total = 0
    for _ in range(n):
        u = min((i for i in range(n) if not in_tree[i]), key=best.__getitem__)
        in_tree[u] = True
        total += best[u]
        ux, uy = points[u]
        for v in range(n):
            if not in_tree[v]:
                d = abs(ux - points[v][0]) + abs(uy - points[v][1])
                best[v] = min(best[v], d)
    return total

print(pipe_cost([(0, 0), (1, 4), (6, 1), (5, 5), (2, 2)]))  # 17

정리

• 문제 1: 다익스트라 한 번, 그다음 최대 거리입니다.
• 문제 2: 추가 자원을 상태에 넣고 층을 쌓은 그래프에서 다익스트라를 돌립니다.
• 문제 3: 비용이 0/1이면 힙이 아니라 덱을 씁니다.
• 문제 4: 완전 그래프에서는 배열 기반 프림 O(n^2)이 크루스칼보다 간단하고 빠릅니다.