BOJ 14621) 나만 안되는 연애

시간 제한 - 2s

메모리 제한 - 256MB

 

---

풀이

모든 노드를 연결할 때 필요한 최소한의 비용을 찾는 문제이다.

다른 점은 노드가 남초대학, 여초대학 둘 중 하나의 값을 가지며, 둘 다 남초 - 남초거나 여초 - 여초라면 연결할 수 없다는 조건이 추가됨

 

입력

입력으로 들어오는 값은
첫 줄에 노드의 갯수 N과 노드 간 간선의 갯수 M

둘째 줄에는 각 학교가 남초 대학('M')인지 여초 대학('W')인지를 나타내는 문자가 N개 입력

셋째 줄부터 M번에 걸쳐 간선의 시작점과 끝점, 거리가 입력됨

 

자료구조

노드의 타입('M' | 'W')을 담기 위한 char형 배열(type)

두 노드와 이를 연결하는 간선의 거리를 담기위한 <int dist, <int start, int end>>형 배열(conn)

 

두 노드의 기존 연결을 확인하기 위해 Union-find에서 사용할 int 배열(parent)

 

로직

1. 노드의 갯수와 간선의 갯수를 입력받아 type과 conn, parent 배열을 초기화 시켜준다.

2. type과 conn을 입력받아준 후 두 노드간의 거리를 기준으로 conn 배열을 정렬시켜준다(오름차순)

3. conn 배열을 순회하며 Union-Find를 통해 Cycle 존재 여부와 두 노드의 타입을 비교해 제약 조건을 우선 확인해준다.

거리를 기준으로 정렬한 배열에서 순차적으로 확인 시, 두 노드가 연결되어 있다면 기존 노드간의 거리가 현재보다 적거나 같음을 만족하기 때문에 Union 과정을 건너뛰어준다.
또한 남초 학교와 여초 학교는 직접적으로 연결될 수 없기 때문에 두 노드의 type이 같다면 Union 과정을 건너뛴다.

 

4. 두 제약 조건을 확인한 이후에는 두 노드를 Union하여 연결해주고 연결에 필요한 거리(비용)를 합해 최소 비용으로 연결된 간선들의 거리를 구해준다.

5. conn 배열의 순회가 끝난 후의 간선의 길이는 cycle이 없음을 보장하나 모든 노드가 연결되었는지 확인을 위해 모든 노드의 Find 연산을 반복해 Root 노드가 동일한지 확인하고 만일 서로 Root노드가 다른 노드가 있다면 두 노드는 연결되지 않은 것으로 간주해 거리의 결과 값을 -1로 무효화시켜준다.

 

코드

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

vector<int> parent;
vector<char> type;
vector<pair<int, pair<int, int>>> conn;

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

void merge(int a, int b) {
	a = find(a);
	b = find(b);
	if (a > b) parent[a] = b;
	else parent[b] = a;
}

int main() {
	int n, m;

	cin >> n >> m;
	
    parent.resize(n + 1);
	conn.resize(m);
	type.resize(n + 1);
    
	for (int i = 0; i <= n; i++) parent[i] = i;
	for (int i = 1; i <= n; i++) cin >> type[i];
	for (int i = 0; i < m; i++) cin >> conn[i].second.first >> conn[i].second.second >> conn[i].first;
	
    sort(conn.begin(), conn.end());

	int ret = 0;
	for (int i = 0; i < m; i++) {	// 필요 간선 연결
		if (type[conn[i].second.first] == type[conn[i].second.second]) continue ;
		if (find(conn[i].second.first) == find(conn[i].second.second)) continue ;
		merge(conn[i].second.first, conn[i].second.second);
		ret += conn[i].first;
	}
	
    //	노드 전체 연결 확인
    for (int i = 1; i < n; i++) if (find(i) != find(i + 1)) ret = -1;
    
	cout << ret << endl;
	return 0;
}

 

M개의 간선을 거리순으로 정렬할 때, M * log M 만큼의 시간이 걸리므로 전체적인 시간 복잡도도 이와 동일하다.