[BOJ] #1197 최소 스패닝 트리

발행: / 수정:

백준 1197번 - 최소 스패닝 트리의 풀이다.

TL;DR

  • 희소 그래프에서는 Kruskal 알고리즘을, 밀집 그래프에서는 Prim 알고리즘을 사용하자.
  • Kruskal 알고리즘 구현 시 path compression 하나만으로 효율이 충분히 높아진다.
    • Union-by-rank, path halving은 유의미한 성능 개선이 없었다.

답안

Kruskal

그래프가 sparse할 것 같아서 우선 Kruskal로 구현해봤다. 답안은 다음과 같다.

/**
 * @brief BOJ No. 1197 "최소 스패닝 트리"
 *
 * @author Minseo Kim <kimminss0@outlook.kr>
 * @details Solve minimum spanning tree using Kruskal's algorithm with
 * union-find
 */

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

using namespace std;

class UnionFind {
private:
  vector<int> parents; // parent's id if non-root; negated rank if root.

public:
  UnionFind(int num_vertices) : parents(num_vertices + 1, 0) {}

  // path halving
  int find(int vertex) {
    // assert vertex > 0
    while (parents[vertex] > 0) {
      if (parents[parents[vertex]] > 0) {
        parents[vertex] = parents[parents[vertex]];
        vertex = parents[vertex];
      } else {
        return parents[vertex];
      }
    }
    return vertex;
  }

  // union
  bool unite(int vertex1, int vertex2) {
    int root1 = find(vertex1);
    int root2 = find(vertex2);

    if (root1 == root2) {
      return false;
    }

    int rank1 = -parents[root1];
    int rank2 = -parents[root2];

    if (rank1 > rank2) {
      parents[root2] = root1;
    } else {
      parents[root1] = root2;
      if (rank1 == rank2) {
        parents[root2]--;
      }
    }
    return true;
  }
};

struct Edge {
  int vertex1;
  int vertex2;
  int weight;
  bool operator<(const Edge &that) const { return this->weight < that.weight; }
};

int main(void) {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  int num_vertices, num_edges;
  cin >> num_vertices >> num_edges;

  UnionFind uf{num_vertices};
  vector<Edge> edges;
  edges.reserve(num_edges);

  for (int i = 0; i < num_edges; i++) {
    int v, u, w;
    cin >> v >> u >> w;
    edges.emplace_back(Edge{v, u, w});
  }
  sort(edges.begin(), edges.end());

  int mst_weight = 0;
  for (auto [u, v, weight] : edges) {
    if (uf.unite(u, v)) {
      mst_weight += weight;
    }
  }
  cout << mst_weight;
  return 0;
}
메모리 시간
3356 KB 32 ms

Kruskal 알고리즘을 구현하기 위해 union-find 자료구조를 먼저 구현했다. 원래 union-find는 makeset 연산을 지원하지만 위 구현에서는 생성자가 그 역할을 대신하고, 생성자 외부에서 makeset 연산이 필요하지 않기에 public 메서드로 노출할 필요가 없다.

Union-by-rank와 path halving으로 최적화했다. 일반적인 path compression은 재귀 호출 시 call stack이 쌓이는 문제를 피하고 싶었다.

그러나, 최적화는 path compression 하나로도 충분했음을 알아냈다. 위 코드에서 path halving 대신 path compression으로 구현을 변경해봤으나 실행시간 및 메모리 사용량 차이가 없었다. union-by-rank 또한 적용하지 않은 것과 차이가 없었다.

Union-by-rank와 path halving이 효과적이지 못했던 이유를 추측하자면, find 연산이 매우 자주 일어나므로 path compression 과정에서 재귀가 깊어지기 어렵고, path compression에 의해 트리의 높이가 매우 낮게 유지되므로 union-by-rank을 하지 않아도 union 이후에 트리의 높이가 유의미하게 증가하지 않기 때문인 듯하다.

Path compression 외 다른 최적화를 적용하지 않은 답안은 이 링크에서 확인할 수 있다.

Prim

희소 그래프에서 Prim과 Kruskal의 성능 차이가 궁금하여 Prim을 추가 구현하였다. Adjacency list와 priority queue로 답안을 작성했다.

#include <algorithm>
#include <iostream>
#include <queue>
#include <utility>
#include <vector>

using namespace std;

int main(void) {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  int V, E;
  cin >> V >> E;

  vector<vector<pair<int, int>>> adj(V + 1);
  for (int i = 0; i < E; i++) {
    int v, u, w;
    cin >> v >> u >> w;
    adj[v].emplace_back(make_pair(u, w));
    adj[u].emplace_back(make_pair(v, w));
  }

  const int INF = 1000001;
  vector<bool> vertex_in_mst(V + 1, false);
  vector<int> vertex_cost(V + 1, INF);

  vector<pair<int, int>> _data;
  _data.reserve(V + 1);
  for (int i = 1; i <= V; i++) {
    _data.emplace_back(make_pair(i, INF));
  }
  priority_queue pq(
      [](const auto &a, const auto &b) -> bool { return a.second > b.second; },
      std::move(_data));

  int mst_weight = -INF;
  int i = V;
  while (i) {
    const auto [v, w] = pq.top();
    pq.pop();
    if (vertex_in_mst[v])
      continue;
    vertex_in_mst[v] = true;
    mst_weight += w;
    i--;
    for (auto [u, w] : adj[v]) {
      if (w < vertex_cost[u]) {
        vertex_cost[u] = w;
        pq.emplace(make_pair(u, w));
      }
    }
  }
  cout << mst_weight;
  return 0;
}
메모리 시간
5684 KB 36 ms

Edit(2024-09-24): 위 구현에서 priority queue에 모든 vertices에 대해 {i, INF}를 삽입하고 시작하는 것을 볼 수 있다. 그러나 시작 정점 하나만 priority queue에 넣고 시작하는 것이 효율적이다. 구체적으로는 시작 정점 v0에 대해 cost를 0으로 설정하여 {0, 0}을 삽입하고, mst_weight 또한 -INF 대신 0으로 초기화할 수 있다.

일반적으로 희소 그래프에서는 Kruskal 알고리즘이 Prim 알고리즘보다 효율적이라고 알려져 있다. 실제로 이 문제에서도 메모리 사용량과 실행 시간 모두 Kruskal 알고리즘에 비해 좋지 못했다. 확실히 희소 그래프에서는 힙의 오버헤드와 같은 요인으로 인해 Prim 알고리즘보다 Kruskal 알고리즘이 더 효율적인 것으로 보인다.