[파이썬] 그래프, 트리 알고리즘 정리

파이썬 그래프, 트리 알고리즘 정리

알고리즘 선택 (최단거리)

• 가중치가 없는 경우 - BFS
• 가중치가 있는 경우
  • 음의 가중치가 있을 때 - 벨만-포드
  • 음의 가중치가 없을 때 - 다익스트라
• 전체 쌍 최단 경로 - 플로이드-워셜 ( n ≤ 100 )

---

1. 벨만-포드 (Bellman-Ford Algorithm)

• 시간 복잡도: O(V * E)
• 모든 간선을 확인하며 최단 거리를 찾기 때문에 느리다.
• 음수 간선이 있어도 최단 거리를 찾을 수 있다.

import sys
input = sys.stdin.readline
INF = int(1e9)

n, m = map(int, input().split())
edges = []
dist = [INF] * (n + 1)

for _ in range(m):
    u, v, w = map(int, input().split())
    edges.append((u, v, w))

def bf(start):
    dist[start] = 0
    for i in range(n):
        for j in range(m):
            cur_node, next_node, cost = edges[j]
            if dist[cur_node] != INF and dist[next_node] > dist[cur_node] + cost:
                dist[next_node] = dist[cur_node] + cost
                if i == n - 1:
                    return 1
    return 0

negative_cycle = bf(1)

if negative_cycle:
    print('-1')
else:
    for i in range(2, n+1):
        if dist[i] == INF:
            print('-1')
        else:
            print(dist[i])

2. 다익스트라 (Dijkstra’s Algorithm)

• 시간 복잡도: O(E * log(E))
• 방문하지 않은 노드 중에서 최단 거리가 가장 짧은 노드를 선택하여 한 단계씩 최단 거리를 구해나간다.
• 우선순위 큐 (heap) 사용

import heapq
import sys
input = sys.stdin.readline
INF = int(1e9)

n, m = map(int, input().split())
start = int(input())
graph = [[] for i in range(n + 1)]
dist = [INF] * (n + 1)

for _ in range(m):
    a, b, c = map(int, input().split())
    graph[a].append((b, c))

def dijkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    dist[start] = 0

    while q:
        weight, cur = heapq.heappop(q)
        if dist[cur] < weight:
            continue
        for g in graph[cur]:
            cost = weight + g[1]
            if cost < dist[g[0]]:
                dist[g[0]] = cost
                heapq.heappush(q, (cost, g[0]))

dijkstra(start)

for i in range(1, n + 1):
    if dist[i] == INF:
        print("INFINITY")
    else:
        print(dist[i])

3. 플로이드-워셜 (Floyd-Warshall Algorithm)

• 시간 복잡도: O(V^3) - 3중 for문
• 공간 복잡도: O(V^2) - 2차원 배열

n, m = map(int, input().split())
INF = int(1e9)
graph = [[0 if i == j else INF for i in range(n + 1)] for j in range(n + 1)]

for _ in range(m):
    a, b, c = map(int, input().split())
    graph[a][b] = c

for k in range(1, n + 1):
    for a in range(1, n + 1):
        for b in range(1, n + 1):
            graph[a][b] = min(graph[a][b], graph[a][k] + graph[k][b])

for a in range(1, n + 1):
    for b in range(1, n + 1):
        if graph[a][b] == INF:
            print("INFINITY", end=" ")
        else:
            print(graph[a][b], end=" ")
    print()

---

최소 신장 트리 (MST, Minimum Spanning Tree)

Spanning Tree란?

• 그래프 내의 모든 정점을 포함하는 트리
• 즉, 그래프에서 일부 간선을 선택한 트리

MST란?

• Spanning Tree 중 사용된 간선의 가중치 합이 최소인 트리
• n개의 정점을 가지는 그래프에 대해 n-1 개의 간선만을 사용해야 함
• Cycle이 있으면 안됨

알고리즘 선택

• 간선이 적은 그래프 (희소 그래프, Sparse graph) - 크루스칼 (간선 위주)
• 간선이 많은 그래프(밀집 그래프, Dense Graph) - 프림 (정점 위주)

---

1. 크루스칼 (Kruskal Algorithm)

• 시간 복잡도: O(V^2)
• 간선을 가중치 기준 오름차순으로 정렬 후 추가
• 간선을 기준으로 추가하기 때문에, 사이클 여부를 확인해야 함
• 유니온-파인드 사용

v, e = map(int, input().split())
parent = list(range(v + 1))

def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

edges = []
total_cost = 0

for _ in range(e):
    a, b, cost = map(int, input().split())
    edges.append((cost, a, b))
edges.sort()

for i in range(e):
    cost, a, b = edges[i]
    if find_parent(parent, a) != find_parent(parent, b):
        union_parent(parent, a, b)
        total_cost += cost

print(total_cost)

참고 - 유니온-파인드 (union-find)

def find_parent(parent, x):
    if parent[x] != x:
        return find_parent(parent, parent[x])
    return x

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

v, e = map(int, input().split())
parent = list(range(v + 1))

for i in range(e):
    a, b = map(int, input().split())
    union_parent(parent, a, b)

2. 프림 (Prim Algorithm)

• 시간 복잡도: O(E * log(V))
• 시작 정점을 기준으로 연결된 간선 중 가중치가 가장 작은 간선과 연결된 정점을 추가 후 반복
• 우선순위 큐 (heap) 사용

import heapq
import collections
import sys
sys.setrecursionlimit(10 ** 6)
input = sys.stdin.readline

n, m = map(int,input().split())
graph = collections.defaultdict(list)
visited = [0] * (n + 1)

for i in range(m):
    u, v, cost = map(int,input().split())
    graph[u].append([cost, u, v])
    graph[v].append([cost, v, u])

# 프림 알고리즘
def prim(graph, start_node):
    visited[start_node] = 1
    candidate = graph[start_node]
    heapq.heapify(candidate)
    mst = []
    total_cost = 0

    while candidate:
        cost, u, v = heapq.heappop(candidate)
        if visited[v] == 0:
            visited[v] = 1
            mst.append((u,v))
            total_cost += cost

            for edge in graph[v]:
                if visited[edge[2]] == 0:
                    heapq.heappush(candidate, edge)

    return total_cost

print(prim(graph,1))

---

위상정렬 (Topological Sort)

• 시간 복잡도: O(V + E)
• 유향 그래프
• 그래프 내부에 Cycle이 없어야 함
• 아래 코드는 bfs

from collections import deque

v, e = map(int, input().split())
indegree = [0] * (v + 1)
graph = [[] for _ in range(v + 1)]
for _ in range(e):
    a, b = map(int, input().split())
    graph[a].append(b)
    indegree[b] += 1

def topological_sort():
    result = []
    q = deque()
    for i in range(1, v + 1):
        if indegree[i] == 0:
            q.append(i)

    while q:
        cur = q.popleft()
        result.append(cur)

        for g in graph[cur]:
            indegree[g] -= 1
            if indegree[g] == 0:
                q.append(g)
    for i in result:
        print(i, end=' ')

topological_sort()