最小全域有向木 (src/Graph/minimum_spanning_aborescence.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2024-09-20 12:04:43+09:00
• Include: #include "src/Graph/minimum_spanning_aborescence.hpp"

関数名	概要	計算量
minimum_spanning_aborescense(g,w,root)	最小全域有向木を求める． 引数は Graph クラス と辺の重みを表すvectorと求めたい木の根 root. 返り値は pair でラッピングした二つを返す． 一つ目は MSA の重みの和． 二つ目は MSA に使う辺の番号を表す　vector<int>.	$O(E\log V)$

Depends on

• Union-Find (src/DataStructure/UnionFind.hpp)
• Union-Find (undo 可能) (src/DataStructure/UnionFind_Undoable.hpp)
• グラフ (src/Graph/Graph.hpp)
• CSR 表現を用いた二次元配列 他 (src/Internal/ListRange.hpp)

Verified with

• test/aoj/2647.test.cpp
• test/aoj/GRL_2_B.test.cpp
• test/yosupo/directedmst.test.cpp

Code

#pragma once
#include <utility>
#include "src/Graph/Graph.hpp"
#include "src/DataStructure/UnionFind.hpp"
#include "src/DataStructure/UnionFind_Undoable.hpp"
// return {total cost, edge ids}
// return {0, {}} if the graph has no spanning aborescence of the root
template <class cost_t> std::pair<cost_t, std::vector<int>> minimum_spanning_aborescence(const Graph &g, std::vector<cost_t> w, int root) {
 const int n= g.vertex_size(), m= g.edge_size();
 assert((int)w.size() == m);
 std::vector<cost_t> lz(m);
 std::vector<std::pair<int, int>> lr(m, {-1, -1}), cyc;
 std::vector<int> top(n, -1), es(n, -1);
 UnionFind uf(n);
 UnionFind_Undoable uf2(n);
 auto upd= [&](int i, cost_t v) { w[i]-= v, lz[i]+= v; };
 auto push= [&](int i) {
  auto [l, r]= lr[i];
  if (l != -1) upd(l, lz[i]);
  if (r != -1) upd(r, lz[i]);
  lz[i]= 0;
 };
 auto merge= [&](auto &&rec, int u, int v) -> int {
  if (u == -1) return v;
  if (v == -1) return u;
  if (w[v] < w[u]) std::swap(u, v);
  auto &[l, r]= lr[u];
  return push(u), r= rec(rec, r, v), std::swap(l, r), u;
 };
 for (int i= m; i--;) {
  auto [s, d]= g[i];
  top[d]= merge(merge, top[d], i);
 }
 cost_t sum= 0;
 for (int i= n; i--;) {
  if (i == root) continue;
  for (int v= i;;) {
   if (top[v] == -1) return {};
   int x= uf2.leader(g[es[v]= top[v]].first);
   if (sum+= w[es[v]], upd(es[v], w[es[v]]); uf.unite(v, x)) break;
   int t= uf2.time();
   for (int r; uf2.unite(v, x); v= r, x= uf2.leader(g[es[x]].first)) top[r= uf2.leader(v)]= merge(merge, top[v], top[x]);
   cyc.emplace_back(es[v], t);
   while (top[v] != -1 && uf2.connected(v, g[top[v]].first)) {
    auto [l, r]= lr[top[v]];
    push(top[v]), top[v]= merge(merge, l, r);
   }
  }
 }
 for (auto it= cyc.rbegin(); it != cyc.rend(); ++it) {
  auto [e, t]= *it;
  int r= uf2.leader(g[e].second);
  uf2.rollback(t);
  int v= uf2.leader(g[es[r]].second);
  es[v]= std::exchange(es[r], e);
 }
 es.erase(es.begin() + root);
 return {sum, es};
}

#line 2 "src/Graph/minimum_spanning_aborescence.hpp"
#include <utility>
#line 2 "src/Internal/ListRange.hpp"
#include <vector>
#include <iostream>
#include <iterator>
#include <type_traits>
#define _LR(name, IT, CT) \
 template <class T> struct name { \
  using Iterator= typename std::vector<T>::IT; \
  Iterator bg, ed; \
  Iterator begin() const { return bg; } \
  Iterator end() const { return ed; } \
  size_t size() const { return std::distance(bg, ed); } \
  CT &operator[](int i) const { return bg[i]; } \
 }
_LR(ListRange, iterator, T);
_LR(ConstListRange, const_iterator, const T);
#undef _LR
template <class T> struct CSRArray {
 std::vector<T> dat;
 std::vector<int> p;
 size_t size() const { return p.size() - 1; }
 ListRange<T> operator[](int i) { return {dat.begin() + p[i], dat.begin() + p[i + 1]}; }
 ConstListRange<T> operator[](int i) const { return {dat.cbegin() + p[i], dat.cbegin() + p[i + 1]}; }
};
template <template <class> class F, class T> std::enable_if_t<std::disjunction_v<std::is_same<F<T>, ListRange<T>>, std::is_same<F<T>, ConstListRange<T>>, std::is_same<F<T>, CSRArray<T>>>, std::ostream &> operator<<(std::ostream &os, const F<T> &r) {
 os << '[';
 for (int _= 0, __= r.size(); _ < __; ++_) os << (_ ? ", " : "") << r[_];
 return os << ']';
}
#line 3 "src/Graph/Graph.hpp"
struct Edge: std::pair<int, int> {
 using std::pair<int, int>::pair;
 Edge &operator--() { return --first, --second, *this; }
 int to(int v) const { return first ^ second ^ v; }
 friend std::istream &operator>>(std::istream &is, Edge &e) { return is >> e.first >> e.second, is; }
};
struct Graph: std::vector<Edge> {
 size_t n;
 Graph(size_t n= 0, size_t m= 0): vector(m), n(n) {}
 size_t vertex_size() const { return n; }
 size_t edge_size() const { return size(); }
 size_t add_vertex() { return n++; }
 size_t add_edge(int s, int d) { return emplace_back(s, d), size() - 1; }
 size_t add_edge(Edge e) { return emplace_back(e), size() - 1; }
#define _ADJ_FOR(a, b) \
 for (auto [u, v]: *this) a; \
 for (size_t i= 0; i < n; ++i) p[i + 1]+= p[i]; \
 for (int i= size(); i--;) { \
  auto [u, v]= (*this)[i]; \
  b; \
 }
#define _ADJ(a, b) \
 vector<int> p(n + 1), c(size() << !dir); \
 if (!dir) { \
  _ADJ_FOR((++p[u], ++p[v]), (c[--p[u]]= a, c[--p[v]]= b)) \
 } else if (dir > 0) { \
  _ADJ_FOR(++p[u], c[--p[u]]= a) \
 } else { \
  _ADJ_FOR(++p[v], c[--p[v]]= b) \
 } \
 return {c, p}
 CSRArray<int> adjacency_vertex(int dir) const { _ADJ(v, u); }
 CSRArray<int> adjacency_edge(int dir) const { _ADJ(i, i); }
#undef _ADJ
#undef _ADJ_FOR
};
#line 3 "src/DataStructure/UnionFind.hpp"
#include <algorithm>
class UnionFind {
 std::vector<int> par;
public:
 UnionFind(int n): par(n, -1) {}
 int leader(int u) { return par[u] < 0 ? u : par[u]= leader(par[u]); }
 bool unite(int u, int v) {
  if ((u= leader(u)) == (v= leader(v))) return false;
  if (par[u] > par[v]) std::swap(u, v);
  return par[u]+= par[v], par[v]= u, true;
 }
 bool connected(int u, int v) { return leader(u) == leader(v); }
 int size(int u) { return -par[leader(u)]; }
};
#line 4 "src/DataStructure/UnionFind_Undoable.hpp"
#include <array>
#include <cassert>
class UnionFind_Undoable {
 std::vector<int> par;
 std::vector<std::array<int, 3>> his;
 int cur;
public:
 UnionFind_Undoable(int n): par(n, -1), his{{-1, -1, 1}}, cur(0) { his.reserve(n + 1); }
 int leader(int u) const { return par[u] < 0 ? u : leader(par[u]); }
 bool unite(int u, int v) {
  if (++cur; (u= leader(u)) == (v= leader(v))) return ++his.back()[2], false;
  if (par[u] > par[v]) std::swap(u, v);
  return his.push_back({v, par[v], 1}), par[u]+= par[v], par[v]= u, true;
 }
 bool connected(int u, int v) const { return leader(u) == leader(v); }
 int size(int u) const { return -par[leader(u)]; }
 int time() const { return cur; }
 void undo() {
  if (assert(cur > 0), --cur; --his.back()[2] == 0) {
   auto [v, p, _]= his.back();
   his.pop_back(), par[par[v]]-= p, par[v]= p;
  }
 }
 void rollback(int t) {
  assert(0 <= t), assert(t <= cur);
  if (t == cur) return;
  for (;;) {
   auto &[u, p, i]= his.back();
   if (cur-= i; cur < t) {
    i= t - cur, cur= t;
    break;
   }
   par[par[u]]-= p, par[u]= p, his.pop_back();
  }
 }
};
#line 6 "src/Graph/minimum_spanning_aborescence.hpp"
// return {total cost, edge ids}
// return {0, {}} if the graph has no spanning aborescence of the root
template <class cost_t> std::pair<cost_t, std::vector<int>> minimum_spanning_aborescence(const Graph &g, std::vector<cost_t> w, int root) {
 const int n= g.vertex_size(), m= g.edge_size();
 assert((int)w.size() == m);
 std::vector<cost_t> lz(m);
 std::vector<std::pair<int, int>> lr(m, {-1, -1}), cyc;
 std::vector<int> top(n, -1), es(n, -1);
 UnionFind uf(n);
 UnionFind_Undoable uf2(n);
 auto upd= [&](int i, cost_t v) { w[i]-= v, lz[i]+= v; };
 auto push= [&](int i) {
  auto [l, r]= lr[i];
  if (l != -1) upd(l, lz[i]);
  if (r != -1) upd(r, lz[i]);
  lz[i]= 0;
 };
 auto merge= [&](auto &&rec, int u, int v) -> int {
  if (u == -1) return v;
  if (v == -1) return u;
  if (w[v] < w[u]) std::swap(u, v);
  auto &[l, r]= lr[u];
  return push(u), r= rec(rec, r, v), std::swap(l, r), u;
 };
 for (int i= m; i--;) {
  auto [s, d]= g[i];
  top[d]= merge(merge, top[d], i);
 }
 cost_t sum= 0;
 for (int i= n; i--;) {
  if (i == root) continue;
  for (int v= i;;) {
   if (top[v] == -1) return {};
   int x= uf2.leader(g[es[v]= top[v]].first);
   if (sum+= w[es[v]], upd(es[v], w[es[v]]); uf.unite(v, x)) break;
   int t= uf2.time();
   for (int r; uf2.unite(v, x); v= r, x= uf2.leader(g[es[x]].first)) top[r= uf2.leader(v)]= merge(merge, top[v], top[x]);
   cyc.emplace_back(es[v], t);
   while (top[v] != -1 && uf2.connected(v, g[top[v]].first)) {
    auto [l, r]= lr[top[v]];
    push(top[v]), top[v]= merge(merge, l, r);
   }
  }
 }
 for (auto it= cyc.rbegin(); it != cyc.rend(); ++it) {
  auto [e, t]= *it;
  int r= uf2.leader(g[e].second);
  uf2.rollback(t);
  int v= uf2.leader(g[es[r]].second);
  es[v]= std::exchange(es[r], e);
 }
 es.erase(es.begin() + root);
 return {sum, es};
}

Back to top page