백준 1939 <중량제한> Python

https://www.acmicpc.net/problem/1939

---

Union-Find 알고리즘을 사용하면 보다 간편하게 풀 수 있는 문제입니다.

 

최소 스패닝 트리(MST) 문제를 풀 때 사용하는 크루스칼 알고리즘과 어떻게 보면 비슷한 문제입니다.

일반적으로 크루스칼 알고리즘을 풀 때 엣지를 코스트 기준으로 오름차순 정렬해 순차적으로 병합해가며 연결여부를 탐색합니다. 하지만 해당 문제에서는 모든 노드의 연결이 아닌 연결되었으면 하는 노드가 정해져 있고 최대로 옮길 수 있는 코스트를 찾는 문제이기 때문에 조금 결이 다르다고 볼 수 있습니다.

 

엣지를 코스트 기준으로 내림차순 정렬하면 코스트가 큰 노드먼저 연결하게 됩니다. 그렇게 순차적으로 병합하다보면 어느 순간 연결되었으면 하는 두 노드가 같은 그룹에 속하는 때가 존재하는데 그 경우가 가장 큰 코스트로 두 노드를 연결할 수 있는 순간이므로 문제를 풀 수 있게됩니다.

---

import sys

input = sys.stdin.readline


# Union-Find
# Union
def Union(group_list, n1, n2):
    # 그룹
    g1 = Find(group_list, n1)
    g2 = Find(group_list, n2)
    # 그룹이 같으면
    if g1 == g2:
        # 병합하지 않음
        return False, group_list

    # 그룹이 작은쪽으로 병합
    if g1 < g2:
        group_list[g2] = g1
    
    else:
        group_list[g1] = g2
    
    return True, group_list


# Find
def Find(group_list, node):
    # 자신의 그룹 찾기
    if group_list[node] != node:
        # 재귀적으로 업데이트
        group_list[node] = Find(group_list, group_list[node])

    return group_list[node]


def solution(N, M, S, E, edges):
    # 엣지 내림차순 정렬
    edges = sorted(edges, key=lambda x: x[-1], reverse=True)
    # 그룹 리스트
    group_list = list(range(N+1))
    # S 와 E 의 그룹이 같을 때 까지 병합
    for A, B, C in edges:
        # 병합
        state, group_list = Union(group_list, A, B)
        # 병합 되었으면
        if state:
            # S 와 E 의 그룹이 같으면
            if Find(group_list, S) == Find(group_list, E):
                # 현재 코스트
                print(C)

                return


# 입력
N, M = map(int, input().strip().split())
edges = [list(map(int, input().strip().split())) for _ in range(M)]
S, E = map(int, input().strip().split())

solution(N, M, S, E, edges)