[SWEA 5249번] 최소 신장 트리 - 파이썬(python)

SW Expert Academy

문제

그래프에서 사이클을 제거하고 모든 노드를 포함하는 트리를 구성할 때, 가중치의 합이 최소가 되도록 만든 경우를 최소신장트리라고 한다.

0번부터 V번까지의 노드와 E개의 간선을 가진 그래프 정보가 주어질 때, 이 그래프로부터 최소신장트리를 구성하는 간선의 가중치를 모두 더해 출력하는 프로그램을 만드시오.

입력

첫 줄에 테스트 케이스의 개수 T가 주어지고, 테스트 케이스 별로 첫 줄에 마지막 노드번호 V와 간선의 개수 E가 주어진다.

다음 줄부터 E개의 줄에 걸쳐 간선의 양 끝 노드 n1, n2, 가중치 w가 차례로 주어진다.

1<=T<=50, 1<=V<=1000, 1<=w<=10, 1<=E<=1000000

출력

각 줄마다 "#T" (T는 테스트 케이스 번호)를 출력한 뒤, 답을 출력한다.

풀이

• 최소 신장 트리는 흔히 접할 수 있는 전형적인 알고리즘 문제이다.
• 최소 신장 트리는 크루스칼 알고리즘과 유니온 파인드를 통해 풀이할 수 있다.
• 가중치가 가장 작은 값을 출력해야 하므로 가중치가 작ㅇㄴ 순으로 정렬하는 것이 필요하다.
• 사이클이 발생하는지 확인하기 위해 각 노드의 부모노드가 무엇인지 알 필요가 있다.
• 따라서 부모를 찾고, 부모가 다르다면 연결된 노드에 따라 부모를 재설정해준다.
• 사이클 없이 모든 노드가 연결되었을 때 가중치를 출력한다.

코드

# 부모찾는 함수
def Find(n):
    # 부모가 자신이 아니라면
    if parent[n] != n:
        # 부모 찾기
        parent[n] = Find(parent[n])
    return parent[n]

# 부모관계 설정하는 함수
# 기본 값은 자기자신으로 되어있으므로 간선을 입력 받을 때마다 관계 설정하기
def union(a, b):
    # 부모 찾기
    a = Find(a)
    b= Find(b)
    # 부모가 작은 쪽을 부모로 두고 값을 변경
    if a<b:
        parent[b]=a
    else:
        parent[a]=b


T = int(input())
for test_case in range(1, T + 1):
    # V: 마지막 노드 번호, E: 간선의 개수
    V, E = map(int, input().split())
    # 간선 정보 리스트
    relation=[0]*(E)
    # 가중치를 기준으로 정렬해야 하므로 가중치를 첫번쨰로 입력한 간선 정보를 입력하기
    for i in range(E):
        # n1, n2: 양 끝 노드 번호 , w: 가중치
        n1, n2, w = map(int, input().split())
        relation[i]=[w, n1, n2]
    relation.sort()
    # 부모 테이블 생성 (기본 값은 자기자신)
    parent = list(range(V+1))
    # 최소값
    mn=0
    # 연결이 모두 끝났을 경우 멈추기 위한 cnt
    # 간선 개수 - 1이라면 모두 연결된 것
    cnt=0
    for i in range(E):
        w, n1, n2 = relation[i]
        # n1과 n2가 부모가 같은지 확인
        # 다르다면 union함수로 n1과 n2를 연결해주기
        if Find(n1) != Find(n2):
            union(n1,n2)
            # 가중치 더하기
            mn+=w
            cnt+=1
        # 모두 연결되었다면 break로 종료
        if cnt>=V:
            break

    print(f'#{test_case} {mn}')