[백준] 5257 timeismoney Python 3

최적의 최소 신장 트리(MST) 찾기: 크루스칼 알고리즘과 삼분 탐색

크루스칼 알고리즘과 삼분 탐색을 이용해 최적의 MST를 찾는 방법을 소개합니다. 이 방법은 특히 시간과 비용의 가중치를 조합하여 최적의 MST를 찾는 데 유용합니다.

문제 설명

주어진 그래프에서 모든 정점을 연결하는 간선들의 집합 중, 가중치의 합이 최소가 되는 집합을 찾아야 합니다.
각 간선은 시간과 비용이라는 두 가지 가중치를 가지며, 우리는 이 두 가중치를 조합하여 최적의 MST를 찾고자 합니다.

 

코드 설명

유니온-파인드 자료구조

유니온-파인드는 서로소 집합을 관리하기 위한 자료구조로, 두 집합을 합치거나 특정 원소가 속한 집합을 찾는 작업을 효율적으로 수행할 수 있습니다.

 

크루스칼 알고리즘

크루스칼 알고리즘은 간선을 가중치 순으로 정렬한 후, 사이클이 생기지 않도록 간선을 하나씩 추가해 MST를 만드는 알고리즘입니다.

 

삼분 탐색을 이용한 최적의 가중치 찾기

삼분 탐색은 함수의 최소값을 찾기 위한 방법으로, 주어진 구간을 세 부분으로 나누어 각 부분의 함수 값을 비교하는 방식입니다.

import sys
import heapq
from collections import defaultdict

# 표준 입력을 읽기 위해 사용
input = sys.stdin.read

class UnionFind:
    # 유니온-파인드 자료구조 초기화
    def __init__(self, size):
        self.parent = list(range(size))  # 각 노드의 부모를 자기 자신으로 초기화
        self.rank = [0] * size  # 각 노드의 랭크를 0으로 초기화

    # 주어진 노드의 루트 찾기 (경로 압축 포함)
    def find(self, u):
        if self.parent[u] != u:
            self.parent[u] = self.find(self.parent[u])  # 경로 압축 수행
        return self.parent[u]

    # 두 노드가 속한 집합 합치기 (랭크에 기반한 합병)
    def union(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)
        if root_u != root_v:  # 루트가 다르면 합집합 수행
            if self.rank[root_u] > self.rank[root_v]:
                self.parent[root_v] = root_u
            elif self.rank[root_u] < self.rank[root_v]:
                self.parent[root_u] = root_v
            else:
                self.parent[root_v] = root_u
                self.rank[root_u] += 1

# 크루스칼 알고리즘을 사용하여 최소 신장 트리(MST)를 찾는 함수
def kruskal(n, edges, weight_func):
    edges.sort(key=weight_func)  # 주어진 가중치 함수에 따라 간선 정렬
    uf = UnionFind(n)
    total_time = 0  # 총 시간 초기화
    total_cost = 0  # 총 비용 초기화
    mst_edges = []  # MST 간선 목록 초기화
    for u, v, t, c in edges:
        if uf.find(u) != uf.find(v):  # 사이클이 생기지 않으면 간선 추가
            uf.union(u, v)
            total_time += t  # 총 시간 갱신
            total_cost += c  # 총 비용 갱신
            mst_edges.append((u, v))  # MST 간선 목록에 추가
    return total_time, total_cost, mst_edges

# 주어진 연결 정보를 바탕으로 문제를 해결하는 함수
def solve(n, m, connections):
    edges = [(u, v, t, c) for u, v, t, c in connections]

    # 가중치 함수 정의
    def weight_func(a, b):
        return lambda edge: edge[2] * a + edge[3] * b

    # 주어진 가중치로 MST를 찾는 함수
    def find_mst_with_weights(a, b):
        return kruskal(n, edges, weight_func(a, b))

    # 삼분 탐색을 사용하여 최적의 가중치 조합 찾기
    def ternary_search(s, e):
        a, b = s
        c, d = e

        left = 0.0
        right = 1.0
        while right - left > 1e-6:  # 삼분 탐색의 종료 조건
            left_third = (2*left + right) / 3  # 왼쪽 1/3 지점 계산
            right_third = (left + 2*right) / 3  # 오른쪽 1/3 지점 계산

            # 왼쪽 1/3 지점의 MST 계산
            left_time, left_cost, _ = find_mst_with_weights(left_third, 1-left_third)
            # 오른쪽 1/3 지점의 MST 계산
            right_time, right_cost, _ = find_mst_with_weights(right_third, 1-right_third)

            # MST의 가중치 값 계산
            left_value = left_time * left_cost
            right_value = right_time * right_cost

            # 가중치 값이 작은 쪽으로 범위를 좁힘
            if left_value < right_value:
                right = right_third
            else:
                left = left_third

        # 최적의 가중치 조합으로 MST 계산
        best_time, best_cost, best_edges = find_mst_with_weights((left + right) / 2, 1 - (left + right) / 2)
        return best_time, best_cost, best_edges

    # 초기 범위로 삼분 탐색 수행
    best_time, best_cost, best_edges = ternary_search((1, 0), (0, 1))
    return best_time, best_cost, best_edges

# 메인 함수
def main():
    data = input().split()
    idx = 0
    n = int(data[idx])  # 노드 수
    idx += 1
    m = int(data[idx])  # 간선 수
    idx += 1

    connections = []
    for _ in range(m):
        u = int(data[idx])
        idx += 1
        v = int(data[idx])
        idx += 1
        t = int(data[idx])
        idx += 1
        c = int(data[idx])
        idx += 1
        connections.append((u, v, t, c))  # 연결 정보 추가

    # 문제 해결 및 결과 출력
    best_time, best_cost, best_edges = solve(n, m, connections)
    sys.stdout.write(f"{best_time} {best_cost}\n")
    for u, v in best_edges:
        sys.stdout.write(f"{u} {v}\n")

if __name__ == "__main__":
    main()

코드 동작 원리

1. 입력 받기: 노드 수와 간선 수, 각 간선의 정보를 입력받습니다.
2. 유니온-파인드 초기화: 각 노드를 초기화합니다.
3. 크루스칼 알고리즘: 간선을 가중치 순으로 정렬한 후, MST를 만듭니다.
4. 삼분 탐색: 최적의 가중치를 찾기 위해 삼분 탐색을 수행합니다.
5. 결과 출력: 최적의 MST를 구성하는 간선들의 총 시간과 비용을 출력합니다.

번외

파이썬을 하면서 최초로 루비문제를 파이썬으로 푸는데 성공하였다

게다가 5257 timeismoney에서 파이썬은 나 혼자 성공하여서 감회가 더 새롭다

앞으로도 더 많은 문제에 도전!