그래프 코드
노트의 개념을 Python 3로 옮긴 코드다. 표준 라이브러리만 쓰고, 파일 끝의 assert 예제가 그대로 실행된다. 복사해 python3 파일이름.py로 돌려 보자.
인접 행렬·인접 리스트
"""그래프 표현: 간선 목록을 인접 리스트·인접 행렬로 바꾸고, 차수를 센다.
시간: 리스트 만들기 O(V + E), 행렬 만들기 O(V² + E). 간선 확인은 행렬 O(1), 리스트 O(차수)
공간: 인접 리스트 O(V + E), 인접 행렬 O(V²)
"""
def adjacency_list(nodes: list[str], edges: list[tuple[str, str]], directed: bool) -> dict[str, list[str]]:
"""점마다 이웃 목록(글자 순). 방향 없으면 양쪽에 넣는다."""
neighbors: dict[str, list[str]] = {node: [] for node in nodes}
for start, end in edges:
neighbors[start].append(end)
if not directed:
neighbors[end].append(start)
return {node: sorted(items) for node, items in neighbors.items()}
def adjacency_matrix(nodes: list[str], edges: list[tuple[str, str]], directed: bool) -> list[list[int]]:
"""matrix[i][j] = 1이면 nodes[i] → nodes[j] 간선이 있다."""
index = {node: position for position, node in enumerate(nodes)}
matrix = [[0] * len(nodes) for _ in nodes]
for start, end in edges:
matrix[index[start]][index[end]] = 1
if not directed:
matrix[index[end]][index[start]] = 1
return matrix
def matrix_to_list(nodes: list[str], matrix: list[list[int]]) -> dict[str, list[str]]:
"""행렬을 리스트로 되돌린다."""
return {nodes[row]: [nodes[column] for column, cell in enumerate(cells) if cell] for row, cells in enumerate(matrix)}
def degrees(nodes: list[str], matrix: list[list[int]]) -> tuple[list[int], list[int]]:
"""(진출 차수 = 행의 합, 진입 차수 = 열의 합)."""
out_degree = [sum(row) for row in matrix]
in_degree = [sum(matrix[row][column] for row in range(len(nodes))) for column in range(len(nodes))]
return out_degree, in_degree
if __name__ == "__main__":
nodes = ["A", "B", "C", "D", "E", "F"]
directed_edges = [("A", "B"), ("A", "C"), ("B", "D"), ("D", "E"), ("E", "B"), ("C", "F"), ("F", "E")]
matrix = adjacency_matrix(nodes, directed_edges, directed=True)
assert matrix[0] == [0, 1, 1, 0, 0, 0]
assert matrix_to_list(nodes, matrix) == adjacency_list(nodes, directed_edges, directed=True)
out_degree, in_degree = degrees(nodes, matrix)
assert out_degree == [2, 1, 1, 1, 1, 1]
assert in_degree == [0, 2, 1, 1, 2, 1]
diamond = ["A", "B", "C", "D", "E"]
undirected_edges = [("A", "B"), ("A", "C"), ("B", "D"), ("C", "D"), ("D", "E")]
symmetric = adjacency_matrix(diamond, undirected_edges, directed=False)
assert all(symmetric[i][j] == symmetric[j][i] for i in range(5) for j in range(5))
out_degree, _ = degrees(diamond, symmetric)
assert out_degree == [2, 2, 2, 3, 1]
# 차수의 합은 간선 수의 두 배
assert sum(out_degree) == 2 * len(undirected_edges)
assert adjacency_list(["X"], [], directed=False) == {"X": []}
print("ok")BFS
"""너비 우선 탐색(BFS): 큐로 가까운 점부터 한 겹씩 넓혀 간선 수로 잰 최단 거리와 경로를 구한다.
시간: O(V + E) — 점마다 큐에 한 번, 간선마다 한두 번
공간: O(V) — 큐·거리·이전 점 표
"""
from collections import deque
def bfs(graph: dict[str, list[str]], start: str) -> tuple[list[str], dict[str, int], dict[str, str]]:
"""(꺼낸 순서, 거리, 이전 점). 방문 표시는 큐에 넣을 때 한다."""
distance = {start: 0}
previous: dict[str, str] = {}
order = []
queue = deque([start])
while queue:
current = queue.popleft()
order.append(current)
for neighbor in sorted(graph[current]):
if neighbor in distance:
continue
distance[neighbor] = distance[current] + 1
previous[neighbor] = current
queue.append(neighbor)
return order, distance, previous
def path_to(previous: dict[str, str], start: str, goal: str) -> list[str]:
"""이전 점 표를 거꾸로 따라가 경로를 복원한다. 닿지 않으면 빈 목록."""
if goal != start and goal not in previous:
return []
path = [goal]
while path[-1] != start:
path.append(previous[path[-1]])
return path[::-1]
def grid_distance(rows: list[str], start: tuple[int, int], goal: tuple[int, int]) -> int | None:
"""'#'이 벽인 격자에서 상하좌우 이동 횟수의 최솟값. 못 가면 None."""
distance = {start: 0}
queue = deque([start])
while queue:
row, column = queue.popleft()
if (row, column) == goal:
return distance[goal]
for next_row, next_column in ((row - 1, column), (row + 1, column), (row, column - 1), (row, column + 1)):
inside = 0 <= next_row < len(rows) and 0 <= next_column < len(rows[0])
if not inside or rows[next_row][next_column] == "#" or (next_row, next_column) in distance:
continue
distance[(next_row, next_column)] = distance[(row, column)] + 1
queue.append((next_row, next_column))
return None
def undirected(edges: list[tuple[str, str]]) -> dict[str, list[str]]:
"""간선 목록을 방향 없는 인접 리스트로."""
graph: dict[str, list[str]] = {}
for first, second in edges:
graph.setdefault(first, []).append(second)
graph.setdefault(second, []).append(first)
return graph
if __name__ == "__main__":
graph = undirected([("A", "B"), ("A", "C"), ("B", "C"), ("B", "D"), ("C", "F"), ("D", "E"), ("D", "F"), ("E", "G"), ("F", "G")])
order, distance, previous = bfs(graph, "D")
assert order == ["D", "B", "E", "F", "A", "C", "G"]
assert distance == {"D": 0, "B": 1, "E": 1, "F": 1, "A": 2, "C": 2, "G": 2}
assert path_to(previous, "D", "A") == ["D", "B", "A"]
maze = ["S..#...", ".#.#.#.", ".#...#.", ".###.#.", ".....#G"]
assert grid_distance(maze, (0, 0), (4, 6)) == 14
assert grid_distance(["S#G"], (0, 0), (0, 2)) is None
lonely = {"X": [], "Y": []}
assert bfs(lonely, "X")[1] == {"X": 0}
assert path_to({}, "X", "Y") == []
print("ok")DFS
"""깊이 우선 탐색(DFS): 발견·완료 시각, 방향 그래프의 사이클(역방향 간선), 연결 요소 수.
시간: O(V + E)
공간: O(V) — 재귀 대신 명시적 스택을 써서 깊은 그래프에서도 스택이 넘치지 않는다
"""
def dfs_times(graph: dict[str, list[str]]) -> tuple[dict[str, int], dict[str, int], list[tuple[str, str]]]:
"""글자 순으로 DFS하며 (발견 시각, 완료 시각, 역방향 간선)을 돌려준다."""
discovered: dict[str, int] = {}
finished: dict[str, int] = {}
back_edges: list[tuple[str, str]] = []
clock = 0
for start in sorted(graph):
if start in discovered:
continue
clock += 1
discovered[start] = clock
# 스택 항목: (점, 아직 안 본 이웃 반복자)
stack = [(start, iter(sorted(graph[start])))]
while stack:
node, neighbors = stack[-1]
neighbor = next(neighbors, None)
if neighbor is None:
stack.pop()
clock += 1
finished[node] = clock
continue
if neighbor not in discovered:
clock += 1
discovered[neighbor] = clock
stack.append((neighbor, iter(sorted(graph[neighbor]))))
elif neighbor not in finished:
# 발견했지만 완료 전 = 스택에 걸린 조상으로 돌아가는 간선
back_edges.append((node, neighbor))
return discovered, finished, back_edges
def count_components(graph: dict[str, list[str]]) -> int:
"""방향 없는 그래프에서 DFS를 새로 시작한 횟수 = 연결 요소 수."""
seen: set[str] = set()
components = 0
for start in sorted(graph):
if start in seen:
continue
components += 1
stack = [start]
seen.add(start)
while stack:
for neighbor in graph[stack.pop()]:
if neighbor not in seen:
seen.add(neighbor)
stack.append(neighbor)
return components
if __name__ == "__main__":
directed = {"A": ["B", "C"], "B": ["D"], "C": ["F"], "D": ["E"], "E": ["B"], "F": ["E"]}
discovered, finished, back_edges = dfs_times(directed)
assert discovered == {"A": 1, "B": 2, "D": 3, "E": 4, "C": 8, "F": 9}
assert finished == {"E": 5, "D": 6, "B": 7, "F": 10, "C": 11, "A": 12}
assert back_edges == [("E", "B")]
dag = {"C": ["E"], "A": ["E", "B"], "E": ["F"], "B": ["F", "D"], "F": ["D"], "D": []}
_, finished, back_edges = dfs_times(dag)
assert back_edges == []
# DAG에서 완료 시각이 늦은 순서는 위상 순서다
order = sorted(finished, key=finished.get, reverse=True)
assert all(order.index(start) < order.index(end) for start in dag for end in dag[start])
components = {
"A": ["B", "C"], "B": ["A", "C"], "C": ["A", "B"],
"D": ["E", "F"], "E": ["D", "G"], "F": ["D", "G"], "G": ["E", "F"],
"H": [],
}
assert count_components(components) == 3
assert count_components({}) == 0
print("ok")위상 정렬
"""위상 정렬(칸 알고리즘): 진입 차수가 0인 점을 큐에서 꺼내 순서에 붙이고, 나가는 간선을 지운다.
시간: O(V + E)
공간: O(V) — 진입 차수 표와 큐
"""
from collections import deque
def kahn(graph: dict[str, list[str]]) -> tuple[list[str], list[str]]:
"""(위상 순서, 사이클 때문에 남은 점). 남은 점이 없으면 DAG다."""
in_degree = {node: 0 for node in graph}
for node in graph:
for neighbor in graph[node]:
in_degree[neighbor] += 1
queue = deque(sorted(node for node, degree in in_degree.items() if degree == 0))
order = []
while queue:
current = queue.popleft()
order.append(current)
for neighbor in sorted(graph[current]):
in_degree[neighbor] -= 1
if in_degree[neighbor] == 0:
queue.append(neighbor)
remaining = sorted(node for node in graph if node not in order)
return order, remaining
def is_topological(graph: dict[str, list[str]], order: list[str]) -> bool:
"""모든 간선이 순서상 앞에서 뒤로 향하는가."""
position = {node: index for index, node in enumerate(order)}
return len(position) == len(graph) and all(position[node] < position[neighbor] for node in graph for neighbor in graph[node])
if __name__ == "__main__":
tasks = {"C": ["E"], "A": ["E", "B"], "E": ["G"], "B": ["G", "D"], "G": ["F"], "D": ["F"], "F": []}
order, remaining = kahn(tasks)
assert order == ["A", "C", "B", "E", "D", "G", "F"]
assert remaining == []
assert is_topological(tasks, order)
dag = {"C": ["E"], "A": ["E", "B"], "E": ["F"], "B": ["F", "D"], "F": ["D"], "D": []}
assert is_topological(dag, ["A", "C", "B", "E", "F", "D"])
assert is_topological(dag, ["A", "B", "C", "E", "F", "D"])
assert not is_topological(dag, ["C", "A", "E", "B", "D", "F"])
cyclic = {"A": ["B", "C"], "B": ["D"], "C": ["F"], "D": ["E"], "E": ["B"], "F": ["E"]}
order, remaining = kahn(cyclic)
assert order == ["A", "C", "F"]
assert remaining == ["B", "D", "E"]
assert kahn({}) == ([], [])
print("ok")다익스트라
"""다익스트라: 음수가 아닌 가중치 그래프에서 한 출발점(또는 여러 출발점)의 최단 거리와 경로.
시간: O((V + E) log V) — 이진 힙(heapq). 낡은 항목은 꺼낼 때 버린다
공간: O(V + E) — 거리 표와 힙
"""
import heapq
import math
def dijkstra(graph: dict[str, list[tuple[str, int]]], sources: list[str]) -> tuple[dict[str, float], dict[str, str], int]:
"""(거리, 이전 점, 버린 낡은 항목 수). sources를 모두 거리 0으로 시작한다."""
for _, edges in graph.items():
if any(weight < 0 for _, weight in edges):
raise ValueError("다익스트라는 음수 간선을 다루지 못한다. 벨먼-포드를 쓴다")
distance: dict[str, float] = {node: math.inf for node in graph}
previous: dict[str, str] = {}
heap = []
for source in sources:
distance[source] = 0
heapq.heappush(heap, (0, source))
done: set[str] = set()
stale = 0
while heap:
current_distance, current = heapq.heappop(heap)
if current in done:
stale += 1
continue
done.add(current)
for neighbor, weight in graph[current]:
candidate = current_distance + weight
if candidate < distance[neighbor]:
distance[neighbor] = candidate
previous[neighbor] = current
heapq.heappush(heap, (candidate, neighbor))
return distance, previous, stale
def path_to(previous: dict[str, str], goal: str) -> list[str]:
"""이전 점을 따라 출발점까지 거슬러 올라간 경로."""
path = [goal]
while path[-1] in previous:
path.append(previous[path[-1]])
return path[::-1]
def undirected(edges: list[tuple[str, str, int]]) -> dict[str, list[tuple[str, int]]]:
"""(점, 점, 가중치) 목록을 방향 없는 인접 리스트로."""
graph: dict[str, list[tuple[str, int]]] = {}
for first, second, weight in edges:
graph.setdefault(first, []).append((second, weight))
graph.setdefault(second, []).append((first, weight))
return graph
if __name__ == "__main__":
roads = undirected([
("S", "A", 7), ("S", "B", 2), ("A", "B", 3), ("A", "C", 4),
("B", "D", 9), ("C", "D", 1), ("C", "T", 6), ("D", "T", 3),
])
distance, previous, stale = dijkstra(roads, ["S"])
assert distance == {"S": 0, "A": 5, "B": 2, "C": 9, "D": 10, "T": 13}
assert path_to(previous, "T") == ["S", "B", "A", "C", "D", "T"]
assert stale == 3
# 소방서 S와 T에서 동시에 시작한다(가상 출발점)
nearest, _, _ = dijkstra(roads, ["S", "T"])
assert nearest == {"S": 0, "A": 5, "B": 2, "C": 4, "D": 3, "T": 0}
unreachable, _, _ = dijkstra({"X": [], "Y": []}, ["X"])
assert unreachable["Y"] == math.inf
try:
dijkstra({"S": [("A", -1)], "A": []}, ["S"])
raise AssertionError("음수 간선이면 ValueError가 나야 한다")
except ValueError:
pass
print("ok")벨먼-포드
"""벨먼-포드: 모든 간선을 V-1바퀴 완화해 음수 간선이 있어도 최단 거리를 구하고, 한 바퀴 더 돌려 음수 사이클을 찾는다.
시간: O(VE) — 한 바퀴 동안 줄어든 것이 없으면 일찍 멈춘다
공간: O(V)
"""
import math
def bellman_ford(nodes: list[str], edges: list[tuple[str, str, float]], source: str) -> tuple[dict[str, float], int, bool]:
"""(거리, 거리가 줄어든 바퀴 수, 음수 사이클 여부). 간선은 주어진 순서로 훑는다."""
distance = {node: math.inf for node in nodes}
distance[source] = 0
rounds = 0
for _ in range(len(nodes) - 1):
changed = False
for start, end, weight in edges:
if distance[start] + weight < distance[end]:
distance[end] = distance[start] + weight
changed = True
if not changed:
break
rounds += 1
# V-1바퀴 뒤에도 줄어드는 간선이 있으면 음수 사이클이다
negative_cycle = any(distance[start] + weight < distance[end] for start, end, weight in edges)
return distance, rounds, negative_cycle
if __name__ == "__main__":
nodes = ["S", "A", "B", "C", "D"]
edges = [("C", "D", 2), ("A", "C", 3), ("B", "A", -4), ("D", "B", 1), ("S", "A", 5), ("S", "B", 6), ("A", "D", 8)]
distance, rounds, negative = bellman_ford(nodes, edges, "S")
assert distance == {"S": 0, "A": 2, "B": 6, "C": 5, "D": 7}
assert rounds == 4
assert not negative
# 같은 간선도 좋은 순서로 훑으면 한 바퀴면 된다
bad_order = [("C", "D", -3), ("A", "C", 4), ("B", "D", 7), ("B", "A", 3), ("S", "A", 6), ("S", "B", 2)]
good_order = [("S", "A", 6), ("S", "B", 2), ("B", "A", 3), ("A", "C", 4), ("B", "D", 7), ("C", "D", -3)]
assert bellman_ford(nodes, bad_order, "S")[1] == 4
assert bellman_ford(nodes, good_order, "S")[1] == 1
assert bellman_ford(nodes, bad_order, "S")[0] == bellman_ford(nodes, good_order, "S")[0]
# 환율 고리: −ln(환율)의 합이 음수면 차익이 남는다
rates = [("USD", "EUR", 0.9), ("EUR", "JPY", 160), ("JPY", "USD", 0.0072)]
currency_edges = [(start, end, -math.log(rate)) for start, end, rate in rates]
assert bellman_ford(["USD", "EUR", "JPY"], currency_edges, "USD")[2]
assert bellman_ford(["S", "X"], [], "S")[0] == {"S": 0, "X": math.inf}
print("ok")플로이드-워셜
"""플로이드-워셜: 경유해도 되는 점 k를 하나씩 늘리며 모든 쌍의 최단 거리를 구한다.
시간: O(V³) — k가 가장 바깥 반복이어야 한다
공간: O(V²) — 거리 행렬
"""
import math
def floyd_warshall(nodes: list[str], edges: list[tuple[str, str, float]], directed: bool = True) -> list[list[float]]:
"""dist[i][j] = nodes[i]에서 nodes[j]까지의 최단 거리."""
index = {node: position for position, node in enumerate(nodes)}
size = len(nodes)
dist = [[0 if row == column else math.inf for column in range(size)] for row in range(size)]
for start, end, weight in edges:
dist[index[start]][index[end]] = min(dist[index[start]][index[end]], weight)
if not directed:
dist[index[end]][index[start]] = min(dist[index[end]][index[start]], weight)
for via in range(size):
for row in range(size):
for column in range(size):
through = dist[row][via] + dist[via][column]
if through < dist[row][column]:
dist[row][column] = through
return dist
def has_negative_cycle(dist: list[list[float]]) -> bool:
"""자기 자신까지의 거리가 음수인 점이 있으면 음수 사이클 위에 있다."""
return any(dist[node][node] < 0 for node in range(len(dist)))
if __name__ == "__main__":
nodes = ["1", "2", "3", "4"]
edges = [("1", "2", 4), ("2", "3", 3), ("3", "4", 1), ("4", "1", 2), ("2", "4", 7)]
assert floyd_warshall(nodes, edges) == [[0, 4, 7, 8], [6, 0, 3, 4], [3, 7, 0, 1], [2, 6, 9, 0]]
# 그래프의 중심: 가장 먼 점까지 거리(이심률)가 가장 작은 점
towns = ["A", "B", "C", "D", "S", "T"]
roads = [("S", "A", 7), ("S", "B", 2), ("A", "B", 3), ("A", "C", 4), ("B", "D", 9), ("C", "D", 1), ("C", "T", 6), ("D", "T", 3)]
dist = floyd_warshall(towns, roads, directed=False)
eccentricity = [max(row) for row in dist]
assert towns[eccentricity.index(min(eccentricity))] == "A"
assert [sum(row) for row in dist] == [25, 31, 25, 27, 39, 39]
negative = floyd_warshall(["A", "B", "C"], [("A", "B", 2), ("B", "C", -4), ("C", "A", 1)])
assert has_negative_cycle(negative)
assert not has_negative_cycle(dist)
print("ok")유니온 파인드
"""유니온 파인드(서로소 집합): 경로 압축과 랭크 기준 합치기로 find·union을 사실상 상수 시간에 한다.
시간: 연산 m번에 O(m·α(n)) — α는 역 아커만 함수로 실제로는 4 이하
공간: O(n) — 부모 배열과 랭크 배열
"""
class DisjointSet:
"""원소 0..size-1. 부모가 자기 자신이면 집합의 대표(뿌리)다."""
def __init__(self, size: int) -> None:
self.parent = list(range(size))
self.rank = [0] * size
def find(self, element: int) -> int:
"""대표를 찾으며 지나온 원소를 모두 뿌리에 바로 붙인다(경로 압축)."""
root = element
while self.parent[root] != root:
root = self.parent[root]
while self.parent[element] != root:
self.parent[element], element = root, self.parent[element]
return root
def union(self, first: int, second: int) -> bool:
"""두 집합을 합친다. 이미 같은 집합이면 False."""
root_first = self.find(first)
root_second = self.find(second)
if root_first == root_second:
return False
# 랭크가 낮은 뿌리를 높은 뿌리 밑에 붙인다. 같으면 둘째를 첫째 밑에 붙이고 랭크를 올린다
if self.rank[root_first] < self.rank[root_second]:
root_first, root_second = root_second, root_first
self.parent[root_second] = root_first
if self.rank[root_first] == self.rank[root_second]:
self.rank[root_first] += 1
return True
def same(self, first: int, second: int) -> bool:
"""같은 집합인가."""
return self.find(first) == self.find(second)
def count(self) -> int:
"""집합 수 = 뿌리 수."""
return sum(1 for element, parent in enumerate(self.parent) if element == parent)
def depth(parent: list[int], element: int) -> int:
"""뿌리까지 올라가는 걸음 수(경로 압축 없이 잰다)."""
steps = 0
while parent[element] != element:
element = parent[element]
steps += 1
return steps
def naive_union(parent: list[int], first: int, second: int) -> None:
"""비교용: 랭크 없이 늘 첫째의 뿌리를 둘째의 뿌리 밑에 붙인다."""
root_first = first
while parent[root_first] != root_first:
root_first = parent[root_first]
root_second = second
while parent[root_second] != root_second:
root_second = parent[root_second]
if root_first != root_second:
parent[root_first] = root_second
if __name__ == "__main__":
sets = DisjointSet(6)
for first, second in [(0, 1), (2, 3), (4, 5), (1, 3), (5, 2)]:
assert sets.union(first, second)
assert sets.parent == [0, 0, 0, 2, 0, 4]
assert sets.find(5) == 0
assert sets.parent == [0, 0, 0, 2, 0, 0]
assert sets.rank == [2, 0, 1, 0, 1, 0]
assert not sets.union(3, 4)
assert sets.count() == 1
given = DisjointSet(8)
given.parent = [0, 0, 1, 3, 3, 4, 6, 6]
assert given.count() == 3
assert given.find(2) == 0
assert given.same(2, 0) and not given.same(5, 7)
given.find(5)
assert given.parent[5] == 3
# 랭크 기준 합치기는 한 줄로 이어 붙여도 높이 1
chain = DisjointSet(8)
for element in range(7):
chain.union(element, element + 1)
assert max(depth(chain.parent, element) for element in range(8)) == 1
# 랭크 없이 늘 첫째 뿌리를 둘째 밑에 붙이면 한 줄(높이 7)이 된다
naive = list(range(8))
for element in range(7):
naive_union(naive, element, element + 1)
assert naive == [1, 2, 3, 4, 5, 6, 7, 7]
assert depth(naive, 0) == 7
print("ok")크루스칼·프림
"""최소 신장 트리: 크루스칼(간선 정렬 + 유니온 파인드)과 프림(우선순위 큐로 한 그루 키우기).
시간: 크루스칼 O(E log E), 프림 O(E log V)
공간: O(V + E)
"""
import heapq
def kruskal(nodes: list[str], edges: list[tuple[str, str, int]]) -> tuple[int, list[tuple[str, str, int]]]:
"""(총 무게, 고른 간선). 가벼운 간선부터 두 끝이 다른 덩어리일 때만 고른다."""
parent = {node: node for node in nodes}
def find(node: str) -> str:
while parent[node] != node:
# 할아버지를 가리키게 해 길을 반씩 줄인다(경로 반감)
parent[node] = parent[parent[node]]
node = parent[node]
return node
chosen = []
for first, second, weight in sorted(edges, key=lambda edge: edge[2]):
root_first, root_second = find(first), find(second)
if root_first == root_second:
continue
parent[root_first] = root_second
chosen.append((first, second, weight))
if len(chosen) == len(nodes) - 1:
break
return sum(weight for _, _, weight in chosen), chosen
def prim(nodes: list[str], edges: list[tuple[str, str, int]], start: str) -> tuple[int, list[str]]:
"""(총 무게, 트리에 들어온 점 순서). 트리 안과 밖을 잇는 가장 가벼운 간선을 고른다."""
neighbors: dict[str, list[tuple[int, str]]] = {node: [] for node in nodes}
for first, second, weight in edges:
neighbors[first].append((weight, second))
neighbors[second].append((weight, first))
in_tree = {start}
order = [start]
heap = list(neighbors[start])
heapq.heapify(heap)
total = 0
while heap and len(in_tree) < len(nodes):
weight, node = heapq.heappop(heap)
# 두 끝이 이미 트리 안이면 낡은 후보다
if node in in_tree:
continue
in_tree.add(node)
order.append(node)
total += weight
for candidate in neighbors[node]:
if candidate[1] not in in_tree:
heapq.heappush(heap, candidate)
return total, order
def clusters(nodes: list[str], edges: list[tuple[str, str, int]], count: int) -> list[list[str]]:
"""MST에서 가장 무거운 간선 count-1개를 지워 count개 무리로 나눈다."""
_, tree = kruskal(nodes, edges)
kept = sorted(tree, key=lambda edge: edge[2])[: len(tree) - (count - 1)]
groups = {node: {node} for node in nodes}
for first, second, _ in kept:
merged = groups[first] | groups[second]
for node in merged:
groups[node] = merged
unique = {frozenset(group) for group in groups.values()}
return sorted(sorted(group) for group in unique)
if __name__ == "__main__":
nodes = ["A", "B", "C", "D", "E", "F"]
edges = [
("A", "B", 6), ("B", "C", 8), ("A", "D", 3), ("B", "E", 5), ("C", "F", 7),
("D", "E", 2), ("E", "F", 9), ("A", "E", 4), ("B", "F", 1),
]
total, chosen = kruskal(nodes, edges)
assert total == 18
assert [(first, second) for first, second, _ in chosen] == [("B", "F"), ("D", "E"), ("A", "D"), ("B", "E"), ("C", "F")]
prim_total, order = prim(nodes, edges, "C")
assert prim_total == 18
assert order == ["C", "F", "B", "E", "D", "A"]
assert clusters(nodes, edges, 3) == [["A", "D", "E"], ["B", "F"], ["C"]]
# 노트의 삼각형: 최단 경로 트리(무게 4)와 MST(무게 3)는 다르다
assert kruskal(["A", "B", "C"], [("A", "B", 2), ("A", "C", 2), ("B", "C", 1)])[0] == 3
assert kruskal(["X"], []) == (0, [])
print("ok")