[Python | 23324] 어려운 모든 정점 쌍 최단 거리

https://www.acmicpc.net/problem/23324

---

문제 설명

연두는 방금 "플로이드 와샬 알고리즘"을 공부했다. 이 알고리즘은 \( N \)개의 정점으로 이루어진 그래프에서, 모든 정점 쌍의 최단 거리를 \( O(N^3) \)에 구해준다.
신이 난 연두는 자신이 좋아하는 그래프를 하나 가져왔다. 이 그래프는 \( N \)개의 정점과 \( M \)개의 양방향 간선으로 이루어진 단순 연결 그래프이며, 정점에는 \( 1, 2, \ldots, n \)으로 번호가 매겨져있다. 또한 딱 하나의 간선에만 1의 가중치가 있고 나머지 간선은 가중치가 0이다.
이제 이 그래프에서 모든 정점 쌍의 최단 거리의 합을 구해보려고 한다. 즉, \( 1 \le i < j \le N \)를 만족하는 모든 \( \frac{N(N-1)}{2} \)개의 쌍 \( (i, j) \)에 대해, \( i \)번 정점과 \( j \)번 정점 간의 최단 거리의 전부 더한 값을 구할 것이다.
연두는 신나서 코드를 짰지만 한참 동안 기다려도 결과가 나오지 않았다. 절망에 빠진 연두는 더 좋은 방법을 생각해냈는데, 바로 대회에 이 문제를 출제하여 여러분들이 답을 대신 구하게 하는 것이다.

 

📌 시간 제한: 1초

---

입력

첫 번째 줄에 정점의 개수 \( N(2 \le N \le 100\,000) \), 간선의 개수 \( M(1 \le M \le 200\,000) \), 정수 \( K(1 \le K < M) \)가 주어진다.
다음 \( M \)개의 줄에 걸쳐 \( u_i \)와 \( v_i \)가 주어진다. 이것은 \( i \)번째 간선은 \( u_i \)와 \( v_i \)를 잇는다는 것을 의미한다. \( (1 \le u_i, v_i \le n, u_i \ne v_i) \)
단순 연결 그래프만 입력으로 주어지며, \( K \)번째 간선의 가중치는 1이고, 나머지 간선의 가중치는 0이다.

---

출력

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

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

---

문제 접근

“모든 정점 쌍 최단 거리”라는 말을 보면 문제 초반에 적혀 있는 플로이드-와샬이 떠오를 것이다. 하지만 이 문제는 정점이 10만 개까지라 N^3 복잡도의 플로이드-와샬(또는 정점마다 다익스트라)로는 절대 해결되지 않도록 입력 크기가 잡혀 있다. 그래서 입력을 보는 순간, 가중치 1인 간선이 딱 하나뿐이라는 ‘가중치 1인 간선이 단 하나뿐’이라는 특수 구조를 활용하는 쪽으로 시선을 돌려야 한다.

 

이 문제는 정점 N (≤ 10⁵) 개, 간선 M (≤ 2·10⁵) 개의 무방향 그래프가 주어지면 가중치가 0인 간선만 여러 개 있고, 딱 하나만 가중치 1인 간선이 있는 그래프에서 “모든 정점 쌍 최단 거리 합”을 구하는 것이 목표이며, 핵심은 “가중치 1인 K번째 간선이 없어도 두 끝점이 0-가중치 경로로 이미 연결돼 있는가?”를 확인하는 일이다.

 

이게 무슨 뜻이냐면, 문제에서 말하는 목표는 모든 정점 쌍 (i, j)(i < j)의 최단 거리를 계산해 합을 출력하는 것인데 이를 직관적으로 이해하면 다음과 같다.

1. 가중치 0인 간선들만 보면, 같은 연결 컴포넌트 내부에서는 어떤 두 정점 사이도 거리가 0이다.
2. K번째 간선 (u, v)이 없어도 두 정점이 이미 0-경로로 이어져 있으면, 그 간선은 “있으나 마나”이고 결과 합은 전부 0일 것.
3. 반대로 (u, v)를 끊으면 그래프가 정확히 두 컴포넌트 A, B로 갈라진다면 (u, v)는 bridge다. 이때 A와 B에 속한 정점 쌍은 간선을 한 번 건너야 하므로 거리 1, 내부 쌍은 여전히 0이 됨. → 합 = |A| × |B|.

bridge 여부만 판정하면 오버헤드 없이 답이 바로 나오므로, 문제를 “(u,v) 이 0-간선들로도 이어져 있나?”로 축소할 수 있다.

 

예제로 이해를 하자면:

 

• 0-간선만 두면 연결 컴포넌트가 {1,2} 과 {3,4,5} 두 개로 쪼개진다.
• 2-3는 두 컴포넌트를 잇는 bridge.
• 결과 합 = |A|×|B| = 2×3 = 6 (서로 다른 컴포넌트 사이 6쌍, 각 거리 1).

 

• 정점 1·2·3이 모두 0-가중치 경로로 이어져 있다. 즉, 그래프는 하나의 컴포넌트로 쪼개지지 않는다.
• 1-간선 (1, 2) 를 제거해도 1과 2는 0-경로(1–3–2)로 도달 가능 → bridge 아님
• 컴포넌트가 하나이므로, 모든 정점 쌍 거리 = 0 ⇒ 합도 0

이 판단은 K번째 간선의 한쪽 끝 ku에서 시작해 0-1 BFS(가중치 0과 1만 있는 다익스트라의 특수형) 를 한 번만 돌리면 끝난다. 0-가중치만 탐색하므로 사실상 큐에 들어가는 간선 비용은 0뿐이라, 일반 다익스트라와 동일한 코드를 쓰되 우선순위 큐가 사실상 BFS 역할을 한다. 탐색 결과로 얻은 dist 배열에서 kv까지의 거리가 0이면 같은 컴포넌트, 무한대(INF)라면 다른 컴포넌트이다. 방문한 정점 수가 곧 |A|, 나머지가 |B|가 되므로 둘의 곱을 출력하면 된다.

 

왜 다익스트라인가?

• 이 문제의 가중치는 모두 0 또는 1, 즉 음수가 없고 1이 존재한다.
• BFS는 “간선 가중치가 전부 0(동일)”일 때만 최단 거리를 보장한다..
• 다익스트라는 “가중치가 0 이상이면 항상 최단 거리를 올바르게 찾는다”는 일반 해법이므로, 가중치 1을 정확히 처리하려면 다익스트라가 필요하다.
• 플로이드-와샬(O(N^3))이나 정점마다 다익스트라(O(N⋅Mlog⁡N))는 입력 크기에서 시간 초과이지만, ‘가중치 1이 단 하나’라는 구조 덕분에 다익스트라를 한 번만 돌려도 답이 나오므로 O((N+M)log⁡N))으로 충분히 통과된다.

 

왜 0-1 BFS 라고도 부르나?

• 다익스트라의 핵심은 “현재까지 거리 값이 가장 작은 정점을 먼저 확정” 하는 것인데, 가중치가 0 또는 1만 있으면 새로 생길 수 있는 거리는 기존 거리 d 또는 d+1 두 가지뿐이다.
• 이렇게 값이 두 단계로만 증가하기 때문에, 원래 필요하던 우선순위 큐(힙) 대신 덱(deque) 앞·뒤에만 넣어도 “가장 작은 거리부터 꺼내기” 조건이 자동으로 유지도 된다.
• 다익스트라의 자료구조만 덱으로 바꾼 특수 구현을 관용적으로 0-1 BFS라고 부른다.
  • heapq를 쓰면 여전히 다익스트라이고 복잡도는 log⁡N이 붙는다.
  • deque를 쓰면 진짜 0-1 BFS가 되어 O(N+M)으로 더 빨라진다.

---

구현

1. K 번째 간선만 제외하고 인접 리스트를 만든다(가중치 0뿐).
2. 한쪽 끝 u에서 0-1 BFS²(다익스트라의 특수형)를 한 번 수행한다.
   • dist[v] == 0 이면 u와 v는 같은 컴포넌트, dist[v] == INF 이면 다른 컴포넌트.
3. v가 dist[v] == 0 이면 bridge 아님 → 답 0.
   아니면 sizeA = dist == 0 인 정점 수, sizeB = N − sizeA, 답 sizeA × sizeB.

import sys, heapq
input = sys.stdin.readline
INF = 10**18

N, M, K = map(int, input().split())

adj = [[] for _ in range(N + 1)]
ku = kv = None          # K번째(=가중치1) 간선의 양 끝

for i in range(1, M + 1):
    u, v = map(int, input().split())
    if i == K:          # 가중치 1인 간선이므로 일단 저장만
        ku, kv = u, v
    else:               # 나머지는 가중치 0
        adj[u].append((v, 0))
        adj[v].append((u, 0))

# 다익스트라 (여기선 0-1 BFS와 동일)
dist = [INF] * (N + 1)
dist[ku] = 0
pq = [(0, ku)]

while pq:
    d, u = heapq.heappop(pq)
    if d != dist[u]:
        continue
    for v, w in adj[u]:
        if dist[v] > d + w:           # w 는 항상 0
            dist[v] = d + w
            heapq.heappush(pq, (dist[v], v))

# K번째 간선이 0-경로로 대체 가능한지 판단
if dist[kv] == 0:                     # 같은 컴포넌트 → 모든 최단거리 합 0
    print(0)
else:
    sizeA = sum(1 for x in dist[1:] if x == 0)   # ku 가 속한 컴포넌트 크기
    sizeB = N - sizeA
    print(sizeA * sizeB)