[백준][C++] 10888: 두 섬간의 이동

문제 링크: https://www.acmicpc.net/problem/10888

 

disjoint set에 추가 작업을 해줘야 되는 문제입니다.

일단 섬이 1부터 N까지 일렬로 늘어서 있는 점, 입력은 i 섬과 i+1 섬을 잇는 다리를 만든다는 점을 유의해주세요.

 

섬에 다리를 놓을 때마다 집합의 크기를 갱신해줘야 합니다. N 크기의 배열 하나 더 만들어서 루트마다 크기를 저장하면 됩니다. merge 할 때마다 업데이트를 해줘야 하고요

그리고 좀 까다로운게 다리를 이을 때마다 아래 내용을 출력해야 합니다.

• pairs: 왕래가 가능한 (i, j) (i<j) 쌍들의 개수
• bridges: 이때 i에서 j까지 최단거리들의 합

 

Fig 1. 답을 계산하는 방법

예를 들어 1-2의 다리를 이었을 때는 (1, 2)만 왕래가 가능합니다 따라서 pairs=1, 경로도 1개 뿐이니 bridges=1입니다.

1-2-3 이렇게 연결됐을 때는 (1,2) (1,3) + (2,3) 이렇게 3개 쌍이 존재하니 pairs=2+1=3, 경로는 1-2가 1, 1-3이 2, 2-3이 1 따라서 bridges=1+2+1=4가 됩니다. 

 

\[ \displaystyle a_n = a_1 + \sum_{k=1}^{n-1} b_k \tag{1} \]

몇 개 적어보면 계차수열의 형태라는 걸 알 수 있습니다. 일반항을 구해봅시다.

참고로 계차수열의 일반항 공식은 (1)과 같습니다. (까먹어서 찾아봤습니다)

 

\begin{align} \displaystyle
pairs_n &= \sum_{k=1}^{n-1} k = \frac{n(n-1)}{2} \tag{2} \\
bridges_n &= \sum_{k=1}^{n-1} \left(\sum_{i=1}^{k} i\right) = \frac{(n-1)n(n+1)}{6} \tag{3}
\end{align}

이렇게 다리가 이어진 섬이 n개 있을 때 pairs와 bridges를 (2), (3)으로 구할 수 있습니다.

 

그런데 다리가 한 개 생길 때마다 모든 섬 집합을 탐색하면서 pairs와 bridges를 계산하면 시간이 너무 오래걸립니다. 그래서 pairs, bridges를 누적해 계산해야 하는데 merge가 될 때마다 이 값을 업데이트 해줘야 합니다.

merge 될 때마다 두 집합의 pairs, bridges 값을 빼고, merge 후 합쳐진 집합의 크기를 이용해 pairs, bridges 값을 계산해서 누적합니다. 매번 누적한 값만 출력해주면 완료입니다.

#include <bits/stdc++.h>
using namespace std;


vector<int> sz;	// 연결된 도시들의 개수
vector<int> parent;

void init(int sz) {
	parent.resize(sz);
	for (int i = 0; i < sz; ++i)
		parent[i] = i;
}

int find(int u) {
	if (u == parent[u]) return u;
	return parent[u] = find(parent[u]);
}

void merge(int u, int v) {
	u = find(u), v = find(v);
	if (u == v) return;
	parent[u] = v;
}


long long calc_pairs(long long x) {
	return x * (x - 1) / 2LL;
}

long long calc_bridges(long long x) {
	return ((x - 1) * x * (x + 1)) / 6LL;
}

int main()
{
	ios_base::sync_with_stdio(false);
	cin.tie(nullptr);
	cout.tie(nullptr);
	
	int n;
	cin >> n;
	init(n);
	
	sz.resize(n);
	for (auto& c : sz)
		c = 1;

	int c1, c2;
	long long pairs = 0;
	long long bridges = 0;
	for (int i = 0; i < n - 1; ++i) {
		cin >> c1;
		--c1;

		c1 = find(c1);
		c2 = find(c1 + 1);

		pairs -= calc_pairs(sz[c1]) + calc_pairs(sz[c2]);
		bridges -= calc_bridges(sz[c1]) + calc_bridges(sz[c2]);
		
		sz[c2] += sz[c1];
		merge(c1, c2);

		pairs += calc_pairs(sz[c2]);
		bridges += calc_bridges(sz[c2]);
		
		cout << pairs << ' ' << bridges << '\n';
	}

	return 0;
}