레드 블루 스패닝 트리

• 기본적인 최소 신장 트리 (Minimum Spanning Tree / MST) 알고리즘만 알면 배경지식은 충분하지만, 발상이 재미있다.
• 블루 엣지와 레드 엣지의 코스트를 잘 세팅하면, 블루 엣지를 최대로 사용한 MST와 최소로 사용한 MST를 각각 찾을 수 있다.
• 블루 엣지를 최소로 사용한 MST의 블루 엣지 갯수가 A개, 최대로 사용했을때의 개수가 B개라고 하자. 당연히 블루엣지의 갯수가 A개보다 적거나 B개보다 많은 MST는 존재할수 없다. 그리고 재미있게도, A≤x≤B 인 x에 대해서 블루엣지를 x개 사용한 MST는 항상 존재한다. 아이디어를 떠올리는 것은 어렵지 않지만, 증명은 조금 까다로울수 있는데, 대충의 증명 아이디어는 https://www.acmicpc.net/board/view/43299 을 참고
• 구현은 MST 알고리즘을를 두번 돌리는 것이 전부. 엣지의 코스트가 0 또는 1만 존재하므로, kruskal 알고리즘에서 정렬 과정을 빼고 0인 엣지들을 먼저 추가하고 1인 엣지들을 추가하는 식으로 수행할수 있다. 이렇게 하면 그냥 O(E*α(V)) = O(V^2*α(V))로 MST를 구할수 있다.
  • teflib.tgraph.minimum_spanning_tree_from_edges 에서는 이러한 기능은 제공하지 않으므로, 그냥 따로 구현했다.

"""Solution code for "BOJ 4792. 레드 블루 스패닝 트리".

- Problem link: https://www.acmicpc.net/problem/4792
- Solution link: http://www.teferi.net/ps/problems/boj/4792

Tags: [MST]
"""

import sys
from teflib import disjointset


def minimum_spanning_tree(node_count, zero_edges, one_edges):
    dsu = disjointset.DisjointSet(node_count)
    component_count, total_cost = node_count, 0
    total_cost = 0
    for edges, cost in ((zero_edges, 0), (one_edges, 1)):
        for u, v in edges:
            try:
                dsu.union(u, v, should_raise=True)
            except ValueError:
                continue
            total_cost += cost
            component_count -= 1
            if component_count == 1:
                return total_cost


def main():
    while True:
        n, m, k = [int(x) for x in sys.stdin.readline().split()]
        if n == 0:
            break
        blue_edges, red_edges = [], []
        for _ in range(m):
            c, f, t = sys.stdin.readline().split()
            if c == 'B':
                blue_edges.append((int(f) - 1, int(t) - 1))
            else:
                red_edges.append((int(f) - 1, int(t) - 1))

        min_blue_count = minimum_spanning_tree(n, red_edges, blue_edges)
        min_red_count = minimum_spanning_tree(n, blue_edges, red_edges)
        max_blue_count = n - 1 - min_red_count
        is_possible = min_blue_count <= k <= max_blue_count
        print('1' if is_possible else '0')


if __name__ == '__main__':
    main()

• Dependency: teflib.disjointset.DisjointSet