[Python] BOJ : 1197 최소 스패닝 트리

요구사항

1. 시간 제한 1초 
   1. N이 10^6이면 O(NlogN) 까지
2. 메모리 128MB
   1. 128 * 10^6 100,000,000 정도
3. 최소 스패닝 트리의 가중치를 출력한다. 

설계

1. V, E 를 입력 받는다. 
2. 정점들의 간선에 대한 가중치를 입력받아 오름차순 정렬한다. 크루스칼 알고리즘과 Union-Find 알고리즘을 사용하기 위해.
3. 자기 자신을 부모로 초기화 한다.
4. 두 정점의 부모가 같지 않을 시 두 정점을 연결한다. 
5. 두 정점의 부모가 같을 시 가중치가 더 작은 정점으로 부모를 교체한다.
6. 4, 5번을 간선의 개수만큼 반복한다. 

구현

V, E = map(int, input().split())

edges = [list(map(int, input().split())) for _ in range(E)]
parent = [i for i in range(V+1)]

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

def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)

    if a < b:
        parent[b] = a
    else:
        parent[a] = b
        print(parent)

def smae_parent(a, b):
    return parent[a] ==  parent[b]

result = 0

for a, b, weight in edges:
    if not smae_parent(a, b):
        union_parent(a, b)
        result += weight

print(weight)

실수

• 테스트 코드로 주어진 입력은 가중치가 제일 작은 거 부터 입력되기 때문에 오류가 나지 않았다. 
• 따라서 가중치의 가장 작은 거부터 오름차순으로 정렬될 수 있게 해주어야 한다. 

edges = sorted([list(map(int, input().split())) for _ in range(E)], key=lambda x: x[2])

설명

V, E = map(int, input().split())

• 정점과 간선을 공백을 기준으로 입력 받는다.

edges = sorted([list(map(int, input().split())) for _ in range(E)], key=lambda x: x[2])

• 사용자로부터 E개의 엣지를 입력받아 각 엣지를 리스트로 저장한 후, 세 번째 요소를 기준으로 정렬하는 코드입니다.

result = 0

for a, b, weight in edges:
    if not smae_parent(a, b):
        union_parent(a, b)
        result += weight

• 가중치의 최소합들을 더하기 위해 result를 초기화
• smae -> 오타 same 
• edges의 각 요소의 값 node1, node2, weight 에 대해 
• 만약 a, b의 부모가 같지 않다면
• 두 정점의 부모를 작은 값으로 바꾼다.
• result에 가중치를 더한다. 

def smae_parent(a, b):
    return parent[a] ==  parent[b]

• parent[a] == parent[b] 가 같다면  True 아니라면 False
• 처음에는 자기 자신이 부모이다. [0,1,2,3]
• 따라서 parent[1] == 1, parent[2] == 2 이기 때문에 Fasle

def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)

    if a < b:
        parent[b] = a
    else:
        parent[a] = b
        print(parent)

• False 일 때, union_parent가 실행
• a = find_parent(a)

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

• find_parent(a)가 들어 왔다. (find_parent(1))
• if parent[1] == 1: return 1을 해주자. 

def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)

    if a < b:
        parent[b] = a
    else:
        parent[a] = b
        print(parent)

 

• a = find_parent(a) -> a = 1
• b = 2
• a < b 이므로 
• parent[2] = 1 이 된다. 

for a, b, weight in edges:
    if not smae_parent(a, b):
        union_parent(a, b)
        result += weight

• result += 1을 더한다. (1 2 1) -> 가중치 1

위 과정을 반복하면 됩니다. 

 

제출 시 런타임에러(RecursionError)

Python 이 정한 최대 재귀 깊이 보다 재귀의 깊이가 더 길어질 때 발생하는 에러

아래처럼 바꿔주자.

재귀는 반복문으로 구현 가능하고, 반복문은 재귀로 구현 가능하다.

 

답안 예시 코드

import sys
input = lambda: sys.stdin.readline().rstrip()

V, E = map(int, input().split())

# 크루스칼 알고리즘
edges = []
for _ in range(E):
    A, B, C = map(int, input().split())
    edges.append((A, B, C))
edges.sort(key=lambda x: x[2])  # 비용으로 정렬

# 유니온-파인드
parent = [i for i in range(V + 1)]

def get_parent(x):
    stack = []
    while parent[x] != x:
        stack.append(x)
        x = parent[x]
    root = x
    # 경로 압축
    while stack:
        parent[stack.pop()] = root
    return root

def union_parent(a, b):
    a = get_parent(a)
    b = get_parent(b)

    if a != b:
        if a < b:  # 작은 쪽이 부모가 된다
            parent[b] = a
        else:
            parent[a] = b

def same_parent(a, b):
    return get_parent(a) == get_parent(b)

answer = 0
for a, b, cost in edges:
    if not same_parent(a, b):
        union_parent(a, b)
        answer += cost

print(answer)

바뀐 부분 설명

def get_parent(x):
    stack = []
    while parent[x] != x:  # 부모가 자기 자신이 아닐 때까지 반복
        stack.append(x)     # 현재 노드를 스택에 추가
        x = parent[x]      # 부모 노드로 이동
    root = x              # 루트 노드를 찾음
    # 경로 압축
    while stack:          # 스택에 저장된 노드들에 대해
        parent[stack.pop()] = root  # 각 노드의 부모를 루트로 설정
    return root           # 루트 노드를 반환

예시

가정해봅시다. 다음과 같은 유니온-파인드 구조가 있다고 합시다:

• parent 배열: [0, 1, 1, 2, 2, 3]
  • 노드 1의 부모는 1 (자기 자신)
  • 노드 2의 부모는 1
  • 노드 3의 부모는 2
  • 노드 4의 부모는 2
  • 노드 5의 부모는 3
• 이 구조를 그래프로 표현하면 다음과 같습니다

   1
  / \
 2   3
/ \
4   5

이때, 노드 5의 루트를 찾고 싶다면 get_parent(5)를 호출합니다.

1. 부모 찾기 과정
   • 처음 호출: x = 5
   • 부모 확인: parent[5] = 3 (부모가 자기 자신이 아님)
   • 스택에 5 추가: stack = [5]
   • 이동: x = 3
2. 두 번째 호출
   • 부모 확인: parent[3] = 2
   • 스택에 3 추가: stack = [5, 3]
   • 이동: x = 2
3. 세 번째 호출
   • 부모 확인: parent[2] = 1
   • 스택에 2 추가: stack = [5, 3, 2]
   • 이동: x = 1
4. 네 번째 호출
   • 부모 확인: parent[1] = 1 (부모가 자기 자신)
   • 루트 찾기 완료: root = 1
5. 경로 압축
   • 스택의 모든 노드에 대해 루트를 설정:
     • stack.pop() → 2, parent[2] = 1
     • stack.pop() → 3, parent[3] = 1
     • stack.pop() → 5, parent[5] = 1
   • 최종적으로 parent 배열은 [0, 1, 1, 1, 2, 1]로 업데이트됨.

도움 받은 블로그

https://velog.io/@yoopark/baekjoon-1197