[BOJ/Python] 9344_도로

문제 출처:https://www.acmicpc.net/problem/9344

9344번: 도로

문제

어떤 나라에는 1부터 N까지 이름 붙여진 N개의 도시가 있다. 한 엔지니어는 모든 도시를 연결하는 도로를 건설하고자 한다. 즉, 모든 도시에 대해 항상 다른 어떤 도시로든 이동할 수 있어야 한다. 이때 여러 개의 도시를 통과할 수도 있다. 그의 팀은 몇 개의 길(도로 후보)을 조사했다. 각각의 길은 두 도시를 양방향으로 잇는다. 길 위에 도로를 지을 때는 특정 비용이 든다. (길이 짧을수록 비용도 싸다.)

이 엔지니어는 교통 시스템을 미리 계획하지 않았다. 그는 그저 선호에 따라 한 개의 길을 선택하고, 도로를 건설하는 일을 모든 도시가 연결될 때까지 반복한다.

지금 엔지니어는 도시 p와 도시 q를 잇는 도로를 건설하고자 한다. 비용을 감축하라는 정부의 압력에 의해, 그는 당신에게 그가 해당 도로를 지어야 하는지 여부를 판단하는 프로그램을 작성할 것을 요구했다. 당신의 프로그램은 그 도로를 지으면서 모든 도시를 연결하는 가장 짧은 도로망을 만들 수 있으면 YES라고 대답해야 한다. 그렇지 않다면, NO를 출력해야 한다.

입력

첫 줄에 테스트 케이스의 개수 T가 주어진다. (T ≤ 10) 

각 테스트 케이스의 첫 줄에는 4개의 정수 N, M, p, q가 주어진다. N(2 ≤ N ≤ 10,000)은 도로망 위에 존재하는 도시의 수이다. M(1 ≤ M ≤ 20,000)은 길의 수이다. p와 q(1 ≤ p,q ≤ N)는 그 사이에 도로를 지어도 되는지 판단해야 하는 두 도시이다.

이어지는 M개의 줄 각각에는 u, v, w가 주어진다.(1 ≤ u ≤ N, 1 ≤ v ≤ N, 1 ≤ w ≤ 400,000) 도시 u와 v를 잇는 양방향 길의 비용이 w라는 것을 의미한다. 도로를 짓는 데 드는 비용은 모두 다르며, 두 도시를 잇는 길은 오직 하나이다. 모든 도시를 잇는 도로망이 최소 한 개 이상 존재한다는 것이 보장된다. 모든 입력은 정수이다.

출력

각 테스트 케이스에 대해, p-q를 지으면서 가장 짧은 도로망을 만들 수 있으면 YES를 출력한다. 아니면 NO를 출력한다.

예제 입력 1

3
2 1 1 2
1 2 4
3 3 2 3
1 2 1
1 3 2
2 3 3
4 5 3 4
1 2 1
1 3 3
3 4 2
1 4 4
4 2 5

예제 출력 1

YES
NO
YES

---

풀이

 

문제를 요약해보자.

 

1. 1부터 N까지 이름 붙여진 N개의 도시가 있다.(N개의 노드)
2. 모든 도시에 대해 항상 다른 어떤 도시로든 이동할 수 있어야 한다. (모든 도시가 연결되어 있어야 한다.)
3. 그러면서 가장 짧은 도로망을 만들어야 한다. (간선들의 비용의 합이 최소가 되도록 하여야 한다.)
4. 이때, 도시 p와 도시 q를 잇는 도로가 도로망에 포함된다면 "YES"를 출력하고, 아니라면 "NO"를 출력하는 문제이다.

코드는 다음과 같다.

 

import sys
 
def find_parent(parent,x):
    if parent[x] != x:
        parent[x] = find_parent(parent,parent[x])
    return parent[x]
 
def union_parent(parent,a,b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
 
    if a<b:
        parent[b] = a
    else:
        parent[a] = b
 
 
t = int(sys.stdin.readline())
for _ in range(t):
    n,m,p,q = map(int,sys.stdin.readline().split())
 
    parent = [i for i in range(n+1)]
    edgs = []
    for _ in range(m):
        u,v,w = map(int,sys.stdin.readline().split())
        edgs.append((w,u,v))
    
    edgs.sort()
    result = "NO"
    l = 0
    for w,u,v in edgs:
        if find_parent(parent,u)!= find_parent(parent,v):
            union_parent(parent,u,v)
            l+=1
            if (u==p and v==q) or (u==q and v==p):
                result = "YES"
        if l == n-1:
            break
    
    print(result)

• 크루스칼 알고리즘을 이용하여 MST를 구한다.
• 여기서 간선의 비용이 작은 것들부터 하나씩 살펴보며 연결할 것인지 아닌지 확인한다.
• 만약 해당 간선을 연결한다면, 그 간선이 p-q를 잇는 간선인지 확인하고 맞다면 result를 "YES"로 바꾼다.
• 최종적으로 result를 출력한다.