[백준 / BOJ] C++ 14621 나만 안되는 연애

14621번: 나만 안되는 연애

 

문제

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

14621번: 나만 안되는 연애

 

 

풀이

최소 스패닝 트리(MST) 응용문제로 남초 대학교와 여초 대학교를 연결하는 도로로만 이루어져 있으므로 간선을 입력받을 때 한 학교는 남초, 나머지 학교는 여초이면 된다. 따라서 배열을 만들어 각 학교가 남초인지 여초인지 저장하고 간선을 입력받을 때 조건문만 추가해 주면 된다. 그리고 모든 학교를 연결할 수 없는 경우 -1을 출력해야 하므로 유니온-파인드의 배열 p의 모든 값이 1인지 확인하면 된다. 즉, 배열의 모든 원소의 합이 정점의 수와 같으면 모든 학교가 이어진 것이다.

 

 

코드

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

int V, E, ans = 0;
vector<pair<int, pair<int, int>>> edge;
vector<int> p;

int find(int x) {
  if (p[x] != x)
    p[x] = find(p[x]);
  return p[x];
}

void merge(int x, int y) {
  x = find(x);
  y = find(y);
  if (x < y)
    p[y] = x;
  else
    p[x] = y;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  cin >> V >> E;
  p.resize(V + 1);
  iota(p.begin(), p.end(), 0);
  // 남초/여초 저장
  vector<char> s(V + 1);
  for (int i = 1; i <= V; i++)
    cin >> s[i];
  while (E--) {
    int u, v, w;
    cin >> u >> v >> w;
    if (s[u] != s[v]) // 서로 다르면 간선 추가
      edge.push_back({w, {u, v}});
  }
  sort(edge.begin(), edge.end());
  for (auto [weight, path] : edge)
    if (find(path.first) != find(path.second)) {
      ans += weight;
      merge(path.first, path.second);
    }
  // 배열의 원소가 모두 루트를 가리키도록
  for (int i = 1; i <= V; i++)
    find(i);
  // 모든 원소의 루트가 1인 경우
  // 즉, 배열의 원소의 합이 정점의 수와 같을 때
  if (accumulate(p.begin(), p.end(), 0) == V)
    cout << ans;
  else
    cout << -1;
  return 0;
}