BOJ 23324 - 어려운 모든 정점 쌍 최단 거리

문제

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

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

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

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

입력

첫 번째 줄에 정점의 개수 $N$ ( $2 \le N \le 100\,000$ ), 간선의 개수 $M$ ( $1 \le M \le 200\,000$ ), 정수 $K$ ( $1 \le K \le 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)의 최댓값을 넘을 수 있음에 주의하자.

풀이

여러 풀이 방법이 있겠지만, 이 문제는 가중치가 1인 간선을 기준으로 두 그룹을 나눠 생각하면 훨씬 단순해진다.

가중치가 1인 간선을 잠깐 없애고 보면, 나머지 간선들은 전부 가중치 0이다. 따라서 같은 그룹 안에서 이동하는 최단 거리는 항상 0이고, 서로 다른 그룹에 있는 정점끼리는 결국 그 가중치 1인 간선을 한 번 지나야 하므로 최단 거리가 1이 된다.

문제는 결국 "어떤 두 정점이 같은 그룹인가, 다른 그룹인가"만 판별하면 된다. 이를 위해 parent[] 배열로 각 정점이 속한 최상위 정점, 즉 연결된 컴포넌트의 대표를 기록했다. 두 개의 그룹으로 나누기 위해 가중치가 1인 K번째 간선은 union에 사용하지 않고 건너뛴다. 나머지 간선만 union-find로 합치면 자연스럽게 두 그룹 정보가 채워진다.

이제 가중치가 1인 간선의 양 끝점을 A, B라고 하자. find(A)와 find(B)가 같다면 사실상 그 간선을 쓰지 않아도 두 정점이 이미 같은 그룹에 연결된 것이므로, 모든 정점 쌍의 최단 거리는 전부 0이다. 반대로 대표가 다르면 진짜로 두 그룹이 나뉘어 있다는 뜻이고, 서로 다른 그룹에 속한 정점 쌍만 거리 1을 갖는다.

따라서 이 경우 답은 "첫 번째 그룹의 정점 수 × 두 번째 그룹의 정점 수"가 된다. 핵심은 두 그룹을 어떻게 나누고 판별할지이며, 그 부분을 union-find로 처리하면 계산 자체는 아주 짧아진다. 또한 곱셈 결과가 int 범위를 넘을 수 있으므로 long을 사용해야 한다.

코드

import java.io.*;
import java.util.*;
 
public class Main {
	
	private static int[] parent;
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		int N = Integer.parseInt(st.nextToken());
		int M = Integer.parseInt(st.nextToken());
		int K = Integer.parseInt(st.nextToken());
		int A = 0;
		int B = 0;
		
		parent = new int[N + 1];
		for (int i = 1; i <= N; i++)
			parent[i] = i;
		
		for (int i = 1; i <= M; i++) {
			st = new StringTokenizer(br.readLine());
			int u = Integer.parseInt(st.nextToken());
			int v = Integer.parseInt(st.nextToken());
			
			if (i == K) {
				A = u;
				B = v;
				continue;
			}
			
			union(u, v);
		}
		
		A = find(A);
		B = find(B);
		Long ans = 0L;
		
		if (A != B) {
			Long aCnt = 0L;
			Long bCnt = 0L;
			for (int i = 1; i <= N; i++) {
				int p = find(parent[i]);
				if (p == A) aCnt++;
				else bCnt++;
			}
			
			ans = aCnt * bCnt;
		}
		
		System.out.println(ans);
	}
	
	private static void union(int u, int v) {
		u = find(u);
		v = find(v);
		if (u != v) 
			if (u < v) parent[v] = u;
			else parent[u] = v;
	}
	
	private static int find(int n) {
		if (parent[n] == n) return n;
		return parent[n] = find(parent[n]);
	}
}

복잡도

시간 복잡도: 간선마다 union-find 연산을 수행하므로 $O((N + M) \alpha(N))$ 이다.
공간 복잡도: 부모 배열을 저장하므로 $O(N)$ 이다.

마무리

플로이드-와샬 같은 전형적인 최단 거리 문제처럼 보이지만, 실제로는 가중치가 1인 간선이 딱 하나라는 조건이 전부를 바꾼다. 그 간선을 제외했을 때 그래프가 어떻게 나뉘는지만 보면, 최단 거리 문제를 그룹 크기 계산 문제로 바꿔서 풀 수 있다.

문제#

입력#

출력#

풀이#

코드#

복잡도#

마무리#

문제

입력

출력

풀이

코드

복잡도

마무리