[ 백준 2463 ] 비용 - Python

문제 보기

Union-Find와 우선순위 큐를 활용하여 구현하였다.

Union-Find는 두 정점의 연결 관계를 파악하는 데 사용하였고 우선순위 큐는 간선을 관리하기 위해서 사용하였다.

 

알고리즘은 아래와 같다.

1. 가중치가 높은 순서대로 간선을 추가한다.

2. 해당 간선의 연결 정점을 A, B라고 하였을 때 만약 A와 B의 부모가 다르다면 A와 연결된 집합의 개수 * B와 연결된 집합의 개수 * (전체 간선의 가중치 합 - 현재까지 추가한 가중치 합)을 누적해서 더한다.

3. 해당 정점과 연결된 집합의 개수를 구하기 위해서 자식의 수를 갖는 리스트를 추가적으로 선언하여 관리하였다.

 

import heapq


def find_parent(x):
    if parent[x] != 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
        child_num[a] += child_num[b]
    else:
        parent[a] = b
        child_num[b] += child_num[a]


def find_child_num(x):
    if parent[x] != x:
        return find_child_num(parent[x])
    else:
        return child_num[x]


if __name__ == "__main__":
    n, m = map(int, input().split())  # 정점, 간선의 수
    parent = [-1] * (n + 1)
    child_num = [1] * (n + 1)  # 자기 자신포함 자식 개수
    for i in range(1, n + 1):
        parent[i] = i  # 자기 자신을 가르키도록
    edge = []
    total = 0  # 가중치 총합
    for _ in range(m):  # 간선 입력
        x, y, w = map(int, input().split())
        heapq.heappush(edge, (-w, x, y))  # 간선을 우선순위 큐(최대)로 관리
        total += w

    answer = 0
    current_total = 0
    for _ in range(m):
        w, x, y = heapq.heappop(edge)
        if find_parent(x) != find_parent(y):
            answer += find_child_num(x) * find_child_num(y) * (total - current_total)
            union_parent(x, y)
        current_total -= w
    if answer > 10 ** 9:
        answer %= 10 ** 9
    print(answer)