[ BOJ ] 23324 : 어려운 모든 정점 쌍 최단 거리 ( GOLD 4 ) / Python

>> 문제 바로가기 (https://www.acmicpc.net/problem/23324)

문제

연두는 방금 "플로이드 와샬 알고리즘"을 공부했다. 이 알고리즘은 N개의 정점으로 이루어진 그래프에서, 모든 정점 쌍의 최단 거리를 O(N^3)에 구해준다. 

신이 난 연두는 자신이 좋아하는 그래프를 하나 가져왔다. 이 그래프는 N개의 정점과 M개의 양방향 간선으로 이루어진 단순 연결 그래프이며, 정점에는 1,2,…,n으로 번호가 매겨져있다. 또한 딱 하나의 간선에만 1의 가중치가 있고 나머지 간선은 가중치가 0이다.

이제 이 그래프에서 모든 정점 쌍의 최단 거리의 합을 구해보려고 한다. 즉, 1 ≤ i < j ≤ N를 만족하는 모든 N(N−1)/2개의 쌍 (i,j)에 대해, i번 정점과 j번 정점간의 최단거리를 전부 더한 값을 구할 것이다.

연두는 신나서 코드를 짰지만 한참 동안 기다려도 결과가 나오지 않았다. 절망에 빠진 연두는 더 좋은 방법을 생각해 냈는데, 바로 대회에 이 문제를 출제하여 여러분들이 답을 대신 구하게 하는 것이다.

입력

첫 번째 줄에 정점의 개수 N (2 ≤ N ≤ 100 000), 간선의 개수 M (1 ≤ M ≤ 200 000), 정수 K (1 ≤ K ≤ M)가 주어진다.

다음 M개의 줄에 걸쳐 u_i와 v_i가 주어진다. 이것은 i번째 간선은 u_i와 v_i를 잇는다는 것을 의미한다. (1 ≤ u_i, v_i ≤ n, u_i ≠ v_i)

단순 연결 그래프만 입력으로 주어지며, K번째 간선의 가중치는 1이고, 나머지 간선의 가중치는 0이다.

출력

모든 정점 쌍의 최단 거리의 합을 출력한다.

출력의 값이 32비트형 정수(C/C++의 int)의 최댓값을 넘을 수 있음에 주의하자.

풀이 과정

가야할 길이 무조건 K번째 간선 (가중치가 1인 간선) 만 있는게 아니라면, 최단 거리는 무조건 0이다.

즉, 가야할 길이 무조건 K번째 간선만 있는 정점의 쌍이 답이 된다. ( 가야할 길이 무조건 가중치가 1인 간선만 있으면 최단 거리는 1이므로, 모든 정점 쌍의 최단 거리 합은 가야할 길이 무조건 가중치가 1인 간선인 정점 쌍들 )

 

이 때, 가중치가 1인 간선을 제외하고 그래프를 구성해 보았다.

각각의 그래프 안에서는 최단거리가 0이지만, (1 <-> 2, 3 <-> 4, 3 <-> 5, 4 <-> 5) 그 외의 쌍들은 모두 연결되어 있지 않으므로 가중치가 1인 간선이 필요하다.

현재 그림에서 가중치가 1인 간선이 필요한 정점 쌍은, 1-2 그래프의 정점 개수 x 3-4-5 그래프의 정점 개수 이다. = 6

따라서 가중치가 1인 간선이 필요한 모든 정점 쌍을 구하기 위해선, 각각의 그래프 정점 개수를 서로 곱한 값이 답이 된다.

 

가중치가 1인 간선만 제외하고 BFS를 통해 각각의 그래프마다 정점 개수를 세어 곱셈을 했다.

가중치가 1인 간선을 제외해도 그래프 개수가 1개 ( 전부 연결됨 ) 라면 답은 0이 됨을 유의하자. ( 모두 최단 거리가 0 )

 

import sys
from collections import deque
input = sys.stdin.readline

n, m, k = map(int, input().rstrip().split())
graph = [[] for _ in range(n+1)]
for i in range(m):
    a, b = map(int, input().rstrip().split())
    if i == k - 1: continue
    graph[a].append(b)
    graph[b].append(a)

vis = [0 for _ in range(n+1)]
res = 1
for i in range(1, n+1):
    if vis[i] == 0:
        vis[i] = 1
        cnt = 1
        queue = deque()
        queue.append(i)
        while queue:
            now = queue.popleft()
            for next in graph[now]:
                if vis[next] == 0:
                    vis[next] = 1
                    cnt += 1
                    queue.append(next)
        
        if cnt == n: res = 0
        else: res *= cnt

print(res)