graph/blackbox/mst.hpp

• View this file on GitHub
• Last update: 2025-05-05 02:10:07+09:00
• Include: #include "graph/blackbox/mst.hpp"
• Link: https://atcoder.jp/contests/pakencamp-2023-day1/tasks/pakencamp_2023_day1_p

Depends on

• ds/unionfind/unionfind.hpp

Code

#include "ds/unionfind/unionfind.hpp"

// Brouvka
// find_unused(v)：unused なうちで、v と最小コストで結べる点を探す。(w,cost)
// pair<int,COST> なければ {-1,*} を返すこと。
// 極大な森を返す
template <typename COST, typename F0, typename F1, typename F2>
vc<tuple<int, int, COST>> blackbox_mst(int N, F0 set_used, F1 set_unused, F2 find_unused) {
  using edge = tuple<int, int, COST>;
  UnionFind uf(N);
  vc<edge> res;
  while (1) {
    bool upd = 0;
    vvc<int> comp(N);
    vc<edge> cand(N, {-1, -1, -1});
    FOR(v, N) comp[uf[v]].eb(v);
    FOR(v, N) if (uf[v] == v) {
      for (auto&& x: comp[v]) { set_used(x); }
      for (auto&& x: comp[v]) {
        auto [y, cost] = find_unused(x);
        if (y == -1) continue;
        auto [a, b, c] = cand[v];
        if (a == -1 || cost < c) { cand[v] = {x, y, cost}; }
      }
      for (auto&& x: comp[v]) { set_unused(x); }
    }
    FOR(v, N) if (uf[v] == v) {
      auto [a, b, c] = cand[v];
      if (a == -1) continue;
      upd = 1;
      if (uf.merge(a, b)) res.eb(a, b, c);
    }
    if (!upd) break;
  }
  return res;
}

// Brouvka
// https://atcoder.jp/contests/pakencamp-2023-day1/tasks/pakencamp_2023_day1_p
// incremental なデータ構造を使う版
// init：探索のためのデータ構造初期化等. 2logN 回程度呼ばれる.
// add：探索対象に追加
// find：(w,wt) or (-1,*)
template <typename COST, typename F0, typename F1, typename F2>
vc<tuple<int, int, COST>> blackbox_mst_incremental(int N, F0 init, F1 add, F2 find) {
  using edge = tuple<int, int, COST>;
  UnionFind uf(N);
  vc<edge> res;
  while (uf.n_comp > 1) {
    bool upd = 0;
    vvc<int> comp(N);
    FOR(v, N) comp[uf[v]].eb(v);
    vc<COST> wt(N, infty<COST>);
    vc<pair<int, int>> who(N, {-1, -1});
    init();
    FOR(i, N) {
      for (auto& v: comp[i]) {
        auto [w, x] = find(v);
        if (w != -1 && chmin(wt[i], x)) who[i] = {v, w};
      }
      for (auto& v: comp[i]) { add(v); }
    }
    init();
    FOR_R(i, N) {
      for (auto& v: comp[i]) {
        auto [w, x] = find(v);
        if (w != -1 && chmin(wt[i], x)) who[i] = {v, w};
      }
      for (auto& v: comp[i]) { add(v); }
    }
    FOR(i, N) {
      auto [a, b] = who[i];
      if (a == -1) continue;
      if (uf.merge(a, b)) {
        upd = true;
        res.eb(a, b, wt[i]);
      }
    }
    if (!upd) break;
  }
  return res;
}

// Prim
// O(N^2)
// https://atcoder.jp/contests/pakencamp-2023-day1/tasks/pakencamp_2023_day1_p
template <typename COST, typename F>
vc<tuple<int, int, COST>> mst_prim(int N, F cost) {
  using edge = tuple<int, int, COST>;
  UnionFind uf(N);
  vc<edge> res;
  vc<ll> wt(N, infty<ll>);
  vc<int> TO(N);
  auto add = [&](int v) -> void {
    FOR(w, N) {
      if (TO[w] != -1 && chmin(wt[w], cost(v, w))) TO[w] = v;
    }
    wt[v] = infty<ll>;
    TO[v] = -1;
  };
  add(0);
  FOR(N - 1) {
    int v = min_element(all(wt)) - wt.begin();
    res.eb(TO[v], v, wt[v]);
    add(v);
  }
  return res;
}

#line 2 "ds/unionfind/unionfind.hpp"

struct UnionFind {
  int n, n_comp;
  vc<int> dat; // par or (-size)
  UnionFind(int n = 0) { build(n); }

  void build(int m) {
    n = m, n_comp = m;
    dat.assign(n, -1);
  }

  void reset() { build(n); }

  int operator[](int x) {
    while (dat[x] >= 0) {
      int pp = dat[dat[x]];
      if (pp < 0) { return dat[x]; }
      x = dat[x] = pp;
    }
    return x;
  }

  ll size(int x) {
    x = (*this)[x];
    return -dat[x];
  }

  bool merge(int x, int y) {
    x = (*this)[x], y = (*this)[y];
    if (x == y) return false;
    if (-dat[x] < -dat[y]) swap(x, y);
    dat[x] += dat[y], dat[y] = x, n_comp--;
    return true;
  }

  vc<int> get_all() {
    vc<int> A(n);
    FOR(i, n) A[i] = (*this)[i];
    return A;
  }
};
#line 2 "graph/blackbox/mst.hpp"

// Brouvka
// find_unused(v)：unused なうちで、v と最小コストで結べる点を探す。(w,cost)
// pair<int,COST> なければ {-1,*} を返すこと。
// 極大な森を返す
template <typename COST, typename F0, typename F1, typename F2>
vc<tuple<int, int, COST>> blackbox_mst(int N, F0 set_used, F1 set_unused, F2 find_unused) {
  using edge = tuple<int, int, COST>;
  UnionFind uf(N);
  vc<edge> res;
  while (1) {
    bool upd = 0;
    vvc<int> comp(N);
    vc<edge> cand(N, {-1, -1, -1});
    FOR(v, N) comp[uf[v]].eb(v);
    FOR(v, N) if (uf[v] == v) {
      for (auto&& x: comp[v]) { set_used(x); }
      for (auto&& x: comp[v]) {
        auto [y, cost] = find_unused(x);
        if (y == -1) continue;
        auto [a, b, c] = cand[v];
        if (a == -1 || cost < c) { cand[v] = {x, y, cost}; }
      }
      for (auto&& x: comp[v]) { set_unused(x); }
    }
    FOR(v, N) if (uf[v] == v) {
      auto [a, b, c] = cand[v];
      if (a == -1) continue;
      upd = 1;
      if (uf.merge(a, b)) res.eb(a, b, c);
    }
    if (!upd) break;
  }
  return res;
}

// Brouvka
// https://atcoder.jp/contests/pakencamp-2023-day1/tasks/pakencamp_2023_day1_p
// incremental なデータ構造を使う版
// init：探索のためのデータ構造初期化等. 2logN 回程度呼ばれる.
// add：探索対象に追加
// find：(w,wt) or (-1,*)
template <typename COST, typename F0, typename F1, typename F2>
vc<tuple<int, int, COST>> blackbox_mst_incremental(int N, F0 init, F1 add, F2 find) {
  using edge = tuple<int, int, COST>;
  UnionFind uf(N);
  vc<edge> res;
  while (uf.n_comp > 1) {
    bool upd = 0;
    vvc<int> comp(N);
    FOR(v, N) comp[uf[v]].eb(v);
    vc<COST> wt(N, infty<COST>);
    vc<pair<int, int>> who(N, {-1, -1});
    init();
    FOR(i, N) {
      for (auto& v: comp[i]) {
        auto [w, x] = find(v);
        if (w != -1 && chmin(wt[i], x)) who[i] = {v, w};
      }
      for (auto& v: comp[i]) { add(v); }
    }
    init();
    FOR_R(i, N) {
      for (auto& v: comp[i]) {
        auto [w, x] = find(v);
        if (w != -1 && chmin(wt[i], x)) who[i] = {v, w};
      }
      for (auto& v: comp[i]) { add(v); }
    }
    FOR(i, N) {
      auto [a, b] = who[i];
      if (a == -1) continue;
      if (uf.merge(a, b)) {
        upd = true;
        res.eb(a, b, wt[i]);
      }
    }
    if (!upd) break;
  }
  return res;
}

// Prim
// O(N^2)
// https://atcoder.jp/contests/pakencamp-2023-day1/tasks/pakencamp_2023_day1_p
template <typename COST, typename F>
vc<tuple<int, int, COST>> mst_prim(int N, F cost) {
  using edge = tuple<int, int, COST>;
  UnionFind uf(N);
  vc<edge> res;
  vc<ll> wt(N, infty<ll>);
  vc<int> TO(N);
  auto add = [&](int v) -> void {
    FOR(w, N) {
      if (TO[w] != -1 && chmin(wt[w], cost(v, w))) TO[w] = v;
    }
    wt[v] = infty<ll>;
    TO[v] = -1;
  };
  add(0);
  FOR(N - 1) {
    int v = min_element(all(wt)) - wt.begin();
    res.eb(TO[v], v, wt[v]);
    add(v);
  }
  return res;
}

Back to top page