중급반 7주차 풀이 - DSU, MST

1. [1717] 집합의 표현 (Gold IV)

Union-Find 구현 문제입니다. 특별한 것도 없습니다.

int par[1000020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

void Main(){
	int n, q; cin >> n >> q;
	for (int i = 0; i <= n; i++){ par[i] = i; }
	while (q--){
		int typ; cin >> typ;
		if (typ == 0){
			int a, b; cin >> a >> b; uni(a, b);
		}
		if (typ == 1){
			int a, b; cin >> a >> b;
			cout << ( fnd(a) == fnd(b) ? "YES" : "NO" ) << endl;
		}
	}
}

2. [1976] 여행 가자 (Gold IV)

연결된 도시들끼리 모두 Union을 해봅시다.

그럼 같은 집합에 속한 도시들끼리는 서로 여행을 갈 수 있고, 그렇지 않으면 서로 여행을 갈 수 없게 됩니다.

그러니 여행 계획에서 다른 집합으로 넘어가야 하는 일이 발생하는지 체크해주면 됩니다.

int par[220];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

int arr[1020];

void Main(){
	int n, m; cin >> n >> m;
	for (int i = 1; i <= n; i++){ par[i] = i; }
	for (int i = 1; i <= n; i++){
		for (int j = 1; j <= n; j++){
			int x; cin >> x; if (x){ uni(i, j); }
		}
	}
	for (int i = 1; i <= m; i++){ cin >> arr[i]; }
	for (int i = 1; i < m; i++){
		if ( fnd(arr[i]) != fnd(arr[i+1]) ){ cout << "NO"; return; }
	}
	cout << "YES"; return;
}

3. [18116] 로봇 조립 (Gold IV)

이번에는 집합의 크기를 추가적으로 저장해야 합니다.

초기 상태에서는 모든 집합의 크기가 1이고, 두 다른 집합을 합치면 새로운 집합의 크기는 두 집합 크기의 합이 되므로 이를 그대로 구현해주면 됩니다.

const int N = 1000000;
int par[1000020], siz[1000020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ int pa = fnd(a), pb = fnd(b);
	if (pa == pb){ return; }
	par[pa] = pb; siz[pb] += siz[pa];
}

void Main(){
	for (int i = 1; i <= N; i++){ par[i] = i; siz[i] = 1; }
	int q; cin >> q; while (q--){
		char typ; cin >> typ;
		if (typ == 'I'){ int a, b; cin >> a >> b; uni(a, b); }
		if (typ == 'Q'){ int a; cin >> a; cout << siz[fnd(a)] << endl; }
	}
}

4. [1197] 최소 스패닝 트리 (Gold IV)

MST 연습 문제입니다. Kruskal로 풀어주면 됩니다.

int par[10020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

void Main(){
	int n, m; cin >> n >> m;
	for (int i = 1; i <= n; i++){ par[i] = i; }
	priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
	while (m--){
		int v, w, x; cin >> v >> w >> x;
		pq.push({ x, {v, w} });
	}
	int ans = 0;
	while (!pq.empty()){
		int x = pq.top().fr;
		int v = pq.top().sc.fr, w = pq.top().sc.sc;
		pq.pop();
		if (fnd(v) == fnd(w)){ continue; }
		uni(v, w); ans += x;
	}
	cout << ans;
}

MST 연습 문제입니다. Prim으로 풀어주면 됩니다.

bool chk[10020];
vector<pi2> adj[10020];

void Main(){
	int n, m; cin >> n >> m;
	while (m--){
		int v, w, x; cin >> v >> w >> x;
		adj[v].push_back({w, x}); adj[w].push_back({v, x});
	}
	priority_queue< pi2, vector<pi2>, greater<pi2> > pq;
	chk[1] = 1; for (pi2 p : adj[1]){ pq.push({p.sc, p.fr}); }
	int ans = 0;
	while (!pq.empty()){
		int d = pq.top().fr, v = pq.top().sc; pq.pop();
		if (chk[v]){ continue; } chk[v] = 1; ans += d;
		for (pi2 p : adj[v]){ pq.push({p.sc, p.fr}); }
	}
	cout << ans;
}

5. [1647] 도시 분할 계획 (Gold IV)

MST를 구하되, 두 마을로 나눌 수 있으니 \( N-2 \)개의 간선만 사용하면 됩니다.

int par[100020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

void Main(){
	int n, m; cin >> n >> m;
	for (int i = 1; i <= n; i++){ par[i] = i; }
	priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
	while (m--){
		int v, w, x; cin >> v >> w >> x;
		pq.push({ x, {v, w} });
	}
	int cnt = n-2, ans = 0;
	while (cnt){
		int x = pq.top().fr;
		int v = pq.top().sc.fr, w = pq.top().sc.sc;
		pq.pop();
		if (fnd(v) == fnd(w)){ continue; }
		uni(v, w); ans += x; cnt -= 1;
	}
	cout << ans;
}

6. [14621] 나만 안되는 연애 (Gold III)

MST를 구하면 되지만, '남초' - '여초'를 연결하는 간선들만 남겨야 합니다.

또한, MST가 정말 모든 정점을 포함하고 있는지도 확인해줘야 합니다.

그러니 그렇게 해주면 됩니다.

 

이 문제에서 '경로'가 실제 두 정점간의 경로와는 관련없음에 주의하세요.

int par[1020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

bool chk[1020];

void Main(){
	int n, m; cin >> n >> m;
	for (int i = 1; i <= n; i++){ char ch; cin >> ch; chk[i] = ch=='M'; }
	for (int i = 1; i <= n; i++){ par[i] = i; }
	priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
	while (m--){
		int v, w, x; cin >> v >> w >> x;
		if (chk[v] == chk[w]){ continue; }
		pq.push({ x, {v, w} }); //cout << v << ' ' << w << " / " << x << endl;
	}
	int cnt = n-1, ans = 0;
	while (!pq.empty() && cnt){
		int x = pq.top().fr;
		int v = pq.top().sc.fr, w = pq.top().sc.sc;
		pq.pop();
		if (fnd(v) == fnd(w)){ continue; }
		uni(v, w); ans += x; cnt -= 1;
	}
	cout << (cnt ? -1 : ans);
}

7. [2887] 행성 터널 (Gold I)

간선의 개수가 \( O(N^2) \)이 되는, 무서운 문제입니다.

하지만 MST에 포함될 수 있는 후보는 "한 축을 기준으로 정렬한 뒤, 인접한 도시들을 잇는 간선"들밖에 되지 않습니다.

 

만약 그렇지 않다고 합시다. 즉, 어느 축으로도 정렬해도 인접하지 않지만 MST에 "포함되어야만 하는" 간선 \( v \leftrightarrow w \)가 존재한다고 합시다.

정의에 의해, 어떤 축으로 정렬했을 때 \( v \)와 \( w \) 사이에 \( x \)라는 정점이 존재해야 합니다.

그럼, \( v \leftrightarrow x \leftrightarrow w \)의 두 간선을 사용할 수 있게 되고, \( v \le x \le w \)일 때 |v-x| + |x-w| = |v-w| \)가 되니, \( v \leftrightarrow w \) 대신 \( v \leftrightarrow x \leftrightarrow w \)의 두 간선을 써줄 수 있게 됩니다.

즉, \( v \leftrightarrow w \)가 항상 MST에 들어가야 한다는 가정에 위배됩니다.

 

한 축을 기준으로 정렬한 뒤, 인접한 도시들을 잇는 간선은 축 하나당 \( O(N) \)개니, 이제 약 \( 3N \)개의 간선으로 MST를 구해주면 됩니다.

#define x fr.fr
#define y sc.fr
#define z sc.sc
#define idx fr.sc

int par[100020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

pl4 pos[100020];
inline ll dis(pl4 p1, pl4 p2){ return min({ abs(p1.x-p2.x), abs(p1.y-p2.y), abs(p1.z-p2.z) }); }

void Main(){
	int n; cin >> n;
	for (int i = 1; i <= n; i++){ cin >> pos[i].x >> pos[i].y >> pos[i].z; pos[i].idx = i; }
	for (int i = 1; i <= n; i++){ par[i] = i; }
	priority_queue< pl3, vector<pl3>, greater<pl3> > pq;
	sort(pos+1, pos+n+1, [](pl4 a, pl4 b){ return a.x < b.x; });
	for (int i = 1; i < n; i++){ pq.push({ dis(pos[i], pos[i+1]), {pos[i].idx, pos[i+1].idx} }); }
	sort(pos+1, pos+n+1, [](pl4 a, pl4 b){ return a.y < b.y; });
	for (int i = 1; i < n; i++){ pq.push({ dis(pos[i], pos[i+1]), {pos[i].idx, pos[i+1].idx} }); }
	sort(pos+1, pos+n+1, [](pl4 a, pl4 b){ return a.z < b.z; });
	for (int i = 1; i < n; i++){ pq.push({ dis(pos[i], pos[i+1]), {pos[i].idx, pos[i+1].idx} }); }
	int cnt = n-1, ans = 0;
	while (!pq.empty() && cnt){
		int d = pq.top().fr;
		int v = pq.top().sc.fr, w = pq.top().sc.sc;
		pq.pop();
		if (fnd(v) == fnd(w)){ continue; }
		uni(v, w); ans += d; cnt -= 1;
	}
	cout << (cnt ? -1 : ans);
}

8. [20040] 사이클 게임 (Gold IV)

MST 구하던대로 해주면 됩니다.

간선을 하나씩 입력받으면서 잇다가 이미 두 정점이 같은 집합에 속해있을 때 답을 출력하고 종료해주면 됩니다.

만약 위 경우가 끝까지 나오지 않는다면 0을 출력하고 종료해야 함에 주의하세요.

int par[500020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

void Main(){
	int n, m; cin >> n >> m;
	for (int i = 0; i < n; i++){ par[i] = i; }
	for (int ans = 1; ans <= m; ans++){
		int v, w; cin >> v >> w;
		if (fnd(v) == fnd(w)){ cout << ans; return; }
		uni(v, w);
	}
	cout << 0;
}

9. [14595] 동방 프로젝트 (Large) (Gold II)

Union 연산을 구간에 한 번에 넣는 문제입니다.

이는 집합의 루트를 맨 오른쪽에 놓는 방식으로 빠르게 해결할 수 있으며, 이 이유는

집합의 루트를 찾으면 그 바로 오른쪽 정점은 항상 다른 집합에 속하기 때문입니다.

int par[1000020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

void Main(){
	int n; cin >> n;
	for (int i = 1; i <= n; i++){ par[i] = i; }
	int q; cin >> q; while (q--){
		int st, ed; cin >> st >> ed;
		int ptr = fnd(st); while (ptr < ed){
			uni(ptr, ptr+1); ptr = fnd(ptr+1);
		}
	}
	int ans = 0;
	for (int i = 1; i <= n; i++){ ans += i==fnd(i); }
	cout << ans;
}

10. [17132] 두더지가 정보섬에 올라온 이유 (Platinum V)

가중치가 큰 간선부터 이으면서 양쪽 집합에 생기는 새로운 경로의 수 (= 두 집합의 크기의 곱)만큼 답에 더해주면 됩니다.

즉, 간선을 이으면서 새로 생기는 모든 경로는 정확히 그 가중치와 동일한 행복도를 지니게 됩니다.

이유는 간단한데, 지금까지 이은 모든 간선은 이 간선보다 가중치가 크거나 같기 때문입니다.

ll ans = 0;

int par[100020], siz[100020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b, ll x){ int pa = fnd(a), pb = fnd(b);
	ans += x*siz[pa]*siz[pb];
	par[pa] = pb; siz[pb] += siz[pa];
}

void Main(){
	int n; cin >> n;
	for (int i = 1; i <= n; i++){ par[i] = i; siz[i] = 1; }
	priority_queue<pi3> pq;
	for (int i = 1; i < n; i++){
		int v, w, x; cin >> v >> w >> x;
		pq.push({ x, {v, w} });
	}
	while (!pq.empty()){
		int x = pq.top().fr;
		int v = pq.top().sc.fr, w = pq.top().sc.sc;
		pq.pop(); uni(v, w, x);
	}
	cout << ans;
}

11. [13306] 트리 (Platinum IV)

쿼리를 주어진 순서대로 처리하기에는 너무 어려운 문제입니다.

하지만 이를 역순으로 처리한다면?

트리의 최종 상태인 모든 정점이 다 나눠져 있는 상태에서 시작해서

0번 쿼리로 두 정점을 잇고, 1번 쿼리로 두 정점이 같은 트리 (= 집합)에 속했는지 확인하는 문제가 됩니다.

 

익숙하지 않나요? 위 문제는 "1. [1717] 집합의 표현"과 정확히 동일한 문제이기 때문입니다!

 

정확히는, 두 정점을 이을 때 (나, 나의 부모)를 잘 찾아서 잇는다는 점이 다르긴 합니다.

근데 이건 그냥 각 정점별로 부모를 저장하면 간단하게 처리됩니다.

 

+ 쿼리를 역순으로 처리하니까 출력도 역순으로 해야 함에 주의하세요.

+ Q는 최대 20만이지만, 입력되는 쿼리의 개수는 최대 40만임에 주의하세요.

int par[200020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

int ptr[200020];
pi3 qry[400020];

void Main(){
	int n, q; cin >> n >> q; q += n-1;
	for (int i = 1; i <= n; i++){ par[i] = i; }
	for (int i = 2; i <= n; i++){ cin >> ptr[i]; }
	for (int i = 1; i <= q; i++){
		int typ; cin >> typ;
		if (typ == 0){ int v; cin >> v; qry[i] = { 0, {v, ptr[v]} }; }
		if (typ == 1){ int v, w; cin >> v >> w; qry[i] = { 1, {v, w} }; }
	}
	stack<bool> ans;
	for (int i = q; i >= 1; i--){
		int typ = qry[i].fr;
		int v = qry[i].sc.fr, w = qry[i].sc.sc;
		if (typ == 0){ uni(v, w); }
		if (typ == 1){ ans.push( fnd(v) == fnd(w) ); }
	}
	while (!ans.empty()){
		bool res = ans.top(); ans.pop();
		cout << (res ? "YES" : "NO") << endl;
	}
}

12. [17490] 일감호에 다리 놓기 (Gold III)

공사 중이 아니라 서로 건너갈 수 있는 강의동을 하나의 집합으로 묶으면,

자명히 하나의 강의동 집합에서 와우도까지 가는데 필요한 돌의 개수는 \( \min S_{i} \)가 됩니다.

그러니 이 문제도 결국 Union-Find를 하면서 각 집합별로 필요한 돌의 개수의 최솟값을 같이 관리해주면 됩니다.

 

\( M \le 1 \)이라서 돌 없이도 이동할 수 있는 경우를 예외처리해야 함에 주의하세요.

int par[1000020], cnt[1000020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ int pa = fnd(a), pb = fnd(b);
	par[pa] = pb; cnt[pb] = min(cnt[pb], cnt[pa]);
}

bool chk[1000020];

void Main(){
	int n, m; ll k; cin >> n >> m >> k;
	if (m <= 1){ cout << "YES"; return; }
	for (int i = 1; i <= n; i++){ par[i] = i; cin >> cnt[i]; }
	for (int i = 1; i <= m; i++){
		int v, w; cin >> v >> w;
		if ((v+1)%n == w%n){ chk[v] = 1; }
		if ((w+1)%n == v%n){ chk[w] = 1; }
	}
	for (int i = 1; i <= n; i++){
		if (!chk[i]){
			if (i != n){ uni(i, i+1); } else{ uni(n, 1); }
		}
	}
	ll ans = 0;
	for (int i = 1; i <= n; i++){
		if (fnd(i) != i){ continue; }
		ans += cnt[i];
	}
	cout << (ans <= k ? "YES" : "NO");
}

13. [1774] 우주신과의 교감 (Gold III)

이미 연결된 ( = 필수적으로 연결해야 하는) 간선들이 있는 MST입니다.

풀이는 간단한데, 어차피 연결할 거 그냥 시작하기 전에 연결해주고 시작해주면 됩니다.

 

간선의 가중치 = 두 정점의 거리 이고, 실수오차는 최대한 피하는 걸 추천하는 사람으로서

최종 답인 ans에 더할 때까지는 가중치를 두 정점의 거리의 제곱으로 저장했습니다.

int par[1020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

inline ll dis(pl2 p1, pl2 p2){
	ll y = p1.fr-p2.fr, x = p1.sc-p2.sc;
	return y*y + x*x;
}

pl2 pos[1020];

void Main(){
	int n, m; cin >> n >> m;
	for (int i = 1; i <= n; i++){ par[i] = i; }
	for (int i = 1; i <= n; i++){ cin >> pos[i].fr >> pos[i].sc; }
	priority_queue< pl3, vector<pl3>, greater<pl3> > pq;
	for (int i = 1; i <= n; i++){
		for (int j = i+1; j <= n; j++){
			pq.push({ dis(pos[i], pos[j]), {i, j} });
		}
	}
	while (m--){ int v, w; cin >> v >> w; uni(v, w); }
	ld ans = 0;
	while (!pq.empty()){
		ll x = pq.top().fr;
		int v = pq.top().sc.fr, w = pq.top().sc.sc;
		pq.pop();
		if (fnd(v) == fnd(w)){ continue; }
		uni(v, w); ans += sqrt((ld)x);
	}
	cout << ans;
}

14. [1045] 도로 (Gold I)

Spanning Tree긴 한데, 우선 순위를 최대화하라고 합니다.

그런데 "앞에 있는 간선의 우선순위가 크면" 뒤의 가중치와 상관없이 더 높은 우선순위를 가지기 때문에, 결국 우선순위가 높은 간선부터 차례대로 보면 됩니다.

즉, 우선순위로 간선 재정렬하고 아무 일 없다는 듯이 MST 돌리면 됩니다.

 

+ 도로를 정확히 M개 가져야 하므로, 두 정점이 같은 집합에 속할 때에도 총 M-(N-1)회는 연결시켜줘야 합니다.

int par[60];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

int ans[60];

void Main(){
	priority_queue< pi2, vector<pi2>, greater<pi2> > pq;
	int n, m; cin >> n >> m; m -= n-1;
	for (int i = 0; i < n; i++){ par[i] = i; }
	for (int i = 0; i < n; i++){
		for (int j = 0; j < n; j++){
			char ch; cin >> ch;
			if (ch == 'Y' && i < j){ pq.push({i, j}); }
		}
	}
	int cnt = n-1;
	while (!pq.empty()){
		int v = pq.top().fr, w = pq.top().sc;
		pq.pop();
		if (fnd(v) == fnd(w)){ if (m){ ans[v] += 1; ans[w] += 1; m -= 1; } }
		else{ uni(v, w); ans[v] += 1; ans[w] += 1; cnt -= 1; }
	}
	if (cnt || m){ cout << -1; return; }
	for (int i = 0; i < n; i++){ cout << ans[i] << ' '; }
}

15. [23034] 조별과제 멈춰! (Platinum V)

MST를 구한 뒤 x와 y를 잇는 경로 중 가중치가 가장 큰 간선을 제거하면 됩니다.

x - y를 미리 연결시키고 MST를 돌린다고 생각하면 편합니다.

 

이렇게만 해도 되는 이유가, 위 과정 때문에 MST에 생기는 변화는 x - y를 잇는 경로의 변화밖에 없기 때문입니다.

그렇지 않은 간선이 있다면, 애초부터 MST를 만들 때에 그 간선을 먼저 볼 것이기 때문입니다.

int par[1020];
int fnd(int a){
	if (par[a] == a){ return a; }
	return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }

vector<pi2> adj[1020];
ll dis[1020][1020];
void dfs(int st, int now, int pre, ll mx){
	dis[st][now] = mx;
	for (pi2 p : adj[now]){
		int nxt = p.fr; ll dst = p.sc;
		if (nxt == pre){ continue; }
		dfs(st, nxt, now, max(mx, dst));
	}
}

void Main(){
	priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
	int n, m; cin >> n >> m;
	for (int i = 1; i <= n; i++){ par[i] = i; }
	while (m--){
		int v, w, x; cin >> v >> w >> x;
		pq.push({ x, {v, w} });
	}
	ll ans = 0;
	while (!pq.empty()){
		int x = pq.top().fr;
		int v = pq.top().sc.fr, w = pq.top().sc.sc;
		pq.pop();
		if (fnd(v) == fnd(w)){ continue; }
		uni(v, w); ans += x;
		adj[v].push_back({w, x}); adj[w].push_back({v, x});
	}
	for (int i = 1; i <= n; i++){ dfs(i, i, -1, 0); }
	int q; cin >> q; while (q--){
		int v, w; cin >> v >> w;
		cout << ans - dis[v][w] << endl;
	}
}