directed_mst.hpp¶
求指定根可达所有点的最小权有向生成树（最小树形图），并恢复入边选择。
Complexity: Time: O(E log V). Space: O(V + E).
AC 记录：directedmst。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(E log V).
/// Space: O(V + E).
#include <cassert>
#include <numeric>
#include <optional>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
namespace directed_mst_internal {
template <class T> struct solver {
struct edge {
T cst{};
int u = -1;
int id = -1;
};
struct heap_node {
heap_node *l = nullptr;
heap_node *r = nullptr;
edge val;
T lz{};
int rk = 1;
};
int n;
int rt;
const std::vector<std::tuple<T, int, int>> &es;
std::vector<heap_node> buf;
int ptr = 0;
solver(int n_, int rt_, const std::vector<std::tuple<T, int, int>> &es_)
: n(n_), rt(rt_), es(es_), buf(es.size()) {}
static int heap_rank(heap_node *nd) { return nd == nullptr ? 0 : nd->rk; }
static heap_node *add(heap_node *nd, const T &dlt) {
if (nd != nullptr) {
nd->val.cst += dlt;
nd->lz += dlt;
}
return nd;
}
static void push(heap_node *nd) {
if (nd == nullptr) {
return;
}
add(nd->l, nd->lz);
add(nd->r, nd->lz);
nd->lz = T{};
}
static heap_node *meld(heap_node *a, heap_node *b) {
if (a == nullptr) {
return b;
}
if (b == nullptr) {
return a;
}
if (std::pair(b->val.cst, b->val.id) < std::pair(a->val.cst, a->val.id)) {
std::swap(a, b);
}
push(a);
a->r = meld(a->r, b);
if (heap_rank(a->l) < heap_rank(a->r)) {
std::swap(a->l, a->r);
}
a->rk = heap_rank(a->r) + 1;
return a;
}
heap_node *make_node(T cst, int u, int id) {
heap_node *nd = &buf[ptr++];
nd->val = {cst, u, id};
return nd;
}
static edge pop(heap_node *&nd) {
push(nd);
edge res = nd->val;
nd = meld(nd->l, nd->r);
return res;
}
std::optional<std::pair<T, std::vector<int>>> run() {
assert(n > 0);
assert(0 <= rt && rt < n);
const int cap = 2 * n;
std::vector<heap_node *> in(cap, nullptr);
for (int id = 0; id < int(es.size()); id++) {
auto [cst, u, to] = es[id];
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to) {
in[to] = meld(in[to], make_node(cst, u, id));
}
}
std::vector<int> dsu(cap);
std::iota(dsu.begin(), dsu.end(), 0);
auto fd = [&](int v) {
int uf = v;
while (dsu[uf] != uf) {
uf = dsu[uf];
}
while (dsu[v] != v) {
int nxt = dsu[v];
dsu[v] = uf;
v = nxt;
}
return uf;
};
std::vector<char> st(cap);
std::vector<edge> sel(cap);
std::vector<int> cp(cap, -1);
std::vector<int> uf(cap);
std::iota(uf.begin(), uf.end(), 0);
st[rt] = 2;
int cnt = n;
for (int s = 0; s < n; s++) {
if (st[s] != 0) {
continue;
}
std::vector<int> pth = {s};
while (true) {
int com = pth.back();
st[com] = 1;
edge bst;
bool ok = false;
while (in[com] != nullptr) {
edge can = pop(in[com]);
if (fd(can.u) != fd(com)) {
bst = can;
ok = true;
break;
}
}
if (!ok) {
return std::nullopt;
}
sel[com] = bst;
int pre = uf[fd(bst.u)];
if (st[pre] == 0) {
pth.push_back(pre);
continue;
}
if (st[pre] == 2) {
break;
}
int ctr = cnt++;
assert(ctr < cap);
while (true) {
int mem = pth.back();
pth.pop_back();
in[ctr] = meld(in[ctr], add(in[mem], -sel[mem].cst));
dsu[fd(mem)] = ctr;
cp[mem] = ctr;
st[mem] = 2;
if (mem == pre) {
break;
}
}
uf[fd(ctr)] = ctr;
pth.push_back(ctr);
}
for (int com : pth) {
st[com] = 2;
}
}
std::vector<char> ext(cnt);
ext[rt] = true;
std::vector<int> eid;
for (int com = cnt - 1; com >= 0; com--) {
if (ext[com]) {
continue;
}
int id = sel[com].id;
if (id == -1) {
return std::nullopt;
}
eid.push_back(id);
int v = std::get<2>(es[id]);
while (v != -1 && !ext[v]) {
ext[v] = true;
v = cp[v];
}
}
if (int(eid.size()) != n - 1) {
return std::nullopt;
}
T cst{};
for (int id : eid) {
cst += std::get<0>(es[id]);
}
return std::pair<T, std::vector<int>>{cst, std::move(eid)};
}
};
} // namespace directed_mst_internal
/// @brief Minimum spanning arborescence rooted at root in O(m log n).
/// Edges are (cost, from, to).
/// @return (cost, edge indices), or nullopt if some vertex is unreachable.
template <class T>
std::optional<std::pair<T, std::vector<int>>>
directed_mst(int n, int rt, const std::vector<std::tuple<T, int, int>> &es) {
return directed_mst_internal::solver<T>(n, rt, es).run();
}
} // namespace noya
#ifndef NOYA_DIRECTED_MST_HPP
#define NOYA_DIRECTED_MST_HPP 1
/// @complexity Time: O(E log V).
/// Space: O(V + E).
#include <cassert>
#include <numeric>
#include <optional>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
namespace directed_mst_internal {
template <class T> struct solver {
struct edge {
T cst{};
int u = -1;
int id = -1;
};
struct heap_node {
heap_node *l = nullptr;
heap_node *r = nullptr;
edge val;
T lz{};
int rk = 1;
};
int n;
int rt;
const std::vector<std::tuple<T, int, int>> &es;
std::vector<heap_node> buf;
int ptr = 0;
solver(int n_, int rt_, const std::vector<std::tuple<T, int, int>> &es_)
: n(n_), rt(rt_), es(es_), buf(es.size()) {}
static int heap_rank(heap_node *nd) { return nd == nullptr ? 0 : nd->rk; }
static heap_node *add(heap_node *nd, const T &dlt) {
if (nd != nullptr) {
nd->val.cst += dlt;
nd->lz += dlt;
}
return nd;
}
static void push(heap_node *nd) {
if (nd == nullptr) {
return;
}
add(nd->l, nd->lz);
add(nd->r, nd->lz);
nd->lz = T{};
}
static heap_node *meld(heap_node *a, heap_node *b) {
if (a == nullptr) {
return b;
}
if (b == nullptr) {
return a;
}
if (std::pair(b->val.cst, b->val.id) < std::pair(a->val.cst, a->val.id)) {
std::swap(a, b);
}
push(a);
a->r = meld(a->r, b);
if (heap_rank(a->l) < heap_rank(a->r)) {
std::swap(a->l, a->r);
}
a->rk = heap_rank(a->r) + 1;
return a;
}
heap_node *make_node(T cst, int u, int id) {
heap_node *nd = &buf[ptr++];
nd->val = {cst, u, id};
return nd;
}
static edge pop(heap_node *&nd) {
push(nd);
edge res = nd->val;
nd = meld(nd->l, nd->r);
return res;
}
std::optional<std::pair<T, std::vector<int>>> run() {
assert(n > 0);
assert(0 <= rt && rt < n);
const int cap = 2 * n;
std::vector<heap_node *> in(cap, nullptr);
for (int id = 0; id < int(es.size()); id++) {
auto [cst, u, to] = es[id];
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to) {
in[to] = meld(in[to], make_node(cst, u, id));
}
}
std::vector<int> dsu(cap);
std::iota(dsu.begin(), dsu.end(), 0);
auto fd = [&](int v) {
int uf = v;
while (dsu[uf] != uf) {
uf = dsu[uf];
}
while (dsu[v] != v) {
int nxt = dsu[v];
dsu[v] = uf;
v = nxt;
}
return uf;
};
std::vector<char> st(cap);
std::vector<edge> sel(cap);
std::vector<int> cp(cap, -1);
std::vector<int> uf(cap);
std::iota(uf.begin(), uf.end(), 0);
st[rt] = 2;
int cnt = n;
for (int s = 0; s < n; s++) {
if (st[s] != 0) {
continue;
}
std::vector<int> pth = {s};
while (true) {
int com = pth.back();
st[com] = 1;
edge bst;
bool ok = false;
while (in[com] != nullptr) {
edge can = pop(in[com]);
if (fd(can.u) != fd(com)) {
bst = can;
ok = true;
break;
}
}
if (!ok) {
return std::nullopt;
}
sel[com] = bst;
int pre = uf[fd(bst.u)];
if (st[pre] == 0) {
pth.push_back(pre);
continue;
}
if (st[pre] == 2) {
break;
}
int ctr = cnt++;
assert(ctr < cap);
while (true) {
int mem = pth.back();
pth.pop_back();
in[ctr] = meld(in[ctr], add(in[mem], -sel[mem].cst));
dsu[fd(mem)] = ctr;
cp[mem] = ctr;
st[mem] = 2;
if (mem == pre) {
break;
}
}
uf[fd(ctr)] = ctr;
pth.push_back(ctr);
}
for (int com : pth) {
st[com] = 2;
}
}
std::vector<char> ext(cnt);
ext[rt] = true;
std::vector<int> eid;
for (int com = cnt - 1; com >= 0; com--) {
if (ext[com]) {
continue;
}
int id = sel[com].id;
if (id == -1) {
return std::nullopt;
}
eid.push_back(id);
int v = std::get<2>(es[id]);
while (v != -1 && !ext[v]) {
ext[v] = true;
v = cp[v];
}
}
if (int(eid.size()) != n - 1) {
return std::nullopt;
}
T cst{};
for (int id : eid) {
cst += std::get<0>(es[id]);
}
return std::pair<T, std::vector<int>>{cst, std::move(eid)};
}
};
} // namespace directed_mst_internal
/// @brief Minimum spanning arborescence rooted at root in O(m log n).
/// Edges are (cost, from, to).
/// @return (cost, edge indices), or nullopt if some vertex is unreachable.
template <class T>
std::optional<std::pair<T, std::vector<int>>>
directed_mst(int n, int rt, const std::vector<std::tuple<T, int, int>> &es) {
return directed_mst_internal::solver<T>(n, rt, es).run();
}
} // namespace noya
#endif // NOYA_DIRECTED_MST_HPP
#include <cassert>
#include <numeric>
#include <optional>
#include <tuple>
#include <utility>
#include <vector>
/// @complexity Time: O(E log V).
/// Space: O(V + E).
namespace noya {
namespace directed_mst_internal {
template <class T> struct solver {
struct edge {
T cst{};
int u = -1;
int id = -1;
};
struct heap_node {
heap_node *l = nullptr;
heap_node *r = nullptr;
edge val;
T lz{};
int rk = 1;
};
int n;
int rt;
const std::vector<std::tuple<T, int, int>> &es;
std::vector<heap_node> buf;
int ptr = 0;
solver(int n_, int rt_, const std::vector<std::tuple<T, int, int>> &es_)
: n(n_), rt(rt_), es(es_), buf(es.size()) {}
static int heap_rank(heap_node *nd) { return nd == nullptr ? 0 : nd->rk; }
static heap_node *add(heap_node *nd, const T &dlt) {
if (nd != nullptr) {
nd->val.cst += dlt;
nd->lz += dlt;
}
return nd;
}
static void push(heap_node *nd) {
if (nd == nullptr) {
return;
}
add(nd->l, nd->lz);
add(nd->r, nd->lz);
nd->lz = T{};
}
static heap_node *meld(heap_node *a, heap_node *b) {
if (a == nullptr) {
return b;
}
if (b == nullptr) {
return a;
}
if (std::pair(b->val.cst, b->val.id) < std::pair(a->val.cst, a->val.id)) {
std::swap(a, b);
}
push(a);
a->r = meld(a->r, b);
if (heap_rank(a->l) < heap_rank(a->r)) {
std::swap(a->l, a->r);
}
a->rk = heap_rank(a->r) + 1;
return a;
}
heap_node *make_node(T cst, int u, int id) {
heap_node *nd = &buf[ptr++];
nd->val = {cst, u, id};
return nd;
}
static edge pop(heap_node *&nd) {
push(nd);
edge res = nd->val;
nd = meld(nd->l, nd->r);
return res;
}
std::optional<std::pair<T, std::vector<int>>> run() {
assert(n > 0);
assert(0 <= rt && rt < n);
const int cap = 2 * n;
std::vector<heap_node *> in(cap, nullptr);
for (int id = 0; id < int(es.size()); id++) {
auto [cst, u, to] = es[id];
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to) {
in[to] = meld(in[to], make_node(cst, u, id));
}
}
std::vector<int> dsu(cap);
std::iota(dsu.begin(), dsu.end(), 0);
auto fd = [&](int v) {
int uf = v;
while (dsu[uf] != uf) {
uf = dsu[uf];
}
while (dsu[v] != v) {
int nxt = dsu[v];
dsu[v] = uf;
v = nxt;
}
return uf;
};
std::vector<char> st(cap);
std::vector<edge> sel(cap);
std::vector<int> cp(cap, -1);
std::vector<int> uf(cap);
std::iota(uf.begin(), uf.end(), 0);
st[rt] = 2;
int cnt = n;
for (int s = 0; s < n; s++) {
if (st[s] != 0) {
continue;
}
std::vector<int> pth = {s};
while (true) {
int com = pth.back();
st[com] = 1;
edge bst;
bool ok = false;
while (in[com] != nullptr) {
edge can = pop(in[com]);
if (fd(can.u) != fd(com)) {
bst = can;
ok = true;
break;
}
}
if (!ok) {
return std::nullopt;
}
sel[com] = bst;
int pre = uf[fd(bst.u)];
if (st[pre] == 0) {
pth.push_back(pre);
continue;
}
if (st[pre] == 2) {
break;
}
int ctr = cnt++;
assert(ctr < cap);
while (true) {
int mem = pth.back();
pth.pop_back();
in[ctr] = meld(in[ctr], add(in[mem], -sel[mem].cst));
dsu[fd(mem)] = ctr;
cp[mem] = ctr;
st[mem] = 2;
if (mem == pre) {
break;
}
}
uf[fd(ctr)] = ctr;
pth.push_back(ctr);
}
for (int com : pth) {
st[com] = 2;
}
}
std::vector<char> ext(cnt);
ext[rt] = true;
std::vector<int> eid;
for (int com = cnt - 1; com >= 0; com--) {
if (ext[com]) {
continue;
}
int id = sel[com].id;
if (id == -1) {
return std::nullopt;
}
eid.push_back(id);
int v = std::get<2>(es[id]);
while (v != -1 && !ext[v]) {
ext[v] = true;
v = cp[v];
}
}
if (int(eid.size()) != n - 1) {
return std::nullopt;
}
T cst{};
for (int id : eid) {
cst += std::get<0>(es[id]);
}
return std::pair<T, std::vector<int>>{cst, std::move(eid)};
}
};
} // namespace directed_mst_internal
/// @brief Minimum spanning arborescence rooted at root in O(m log n).
/// Edges are (cost, from, to).
/// @return (cost, edge indices), or nullopt if some vertex is unreachable.
template <class T>
std::optional<std::pair<T, std::vector<int>>>
directed_mst(int n, int rt, const std::vector<std::tuple<T, int, int>> &es) {
return directed_mst_internal::solver<T>(n, rt, es).run();
}
} // namespace noya