directed_mst.hpp¶
平行边归一化、批量左偏堆与缩环求最小树形图，返回总权值和原始边编号。
Complexity: Time: \(O(n + m + M \log (n + 1))\), \(M\) distinct eligible edge pairs. Space: \(O(n + m)\).
AC 记录：directedmst。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(n + m + M log(n + 1)), M distinct eligible edge pairs.
/// Space: O(n + m).
#include <cassert>
#include <cstddef>
#include <limits>
#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;
bool dirty = false;
};
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_) {}
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;
nd->dirty = true;
}
return nd;
}
static void push(heap_node *nd) {
if (nd == nullptr || !nd->dirty) {
return;
}
add(nd->l, nd->lz);
add(nd->r, nd->lz);
nd->lz = T{};
nd->dirty = false;
}
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(0 < n && n <= std::numeric_limits<int>::max() / 2);
assert(0 <= rt && rt < n);
assert(es.size() <= std::size_t(std::numeric_limits<int>::max()));
const int cap = 2 * n;
std::vector<std::size_t> start(std::size_t(n) + 1);
for (const auto &[cst, u, to] : es) {
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to && to != rt) {
++start[to];
}
}
for (int v = 0; v < n; ++v) {
if (v != rt && start[v] == 0) {
return std::nullopt;
}
}
if (n == 1) {
return std::pair<T, std::vector<int>>{T{}, {}};
}
for (int v = 0; v < n; ++v) {
start[v + 1] += start[v];
}
std::vector<int> ids(start.back());
// Stable incoming-edge rows, without a second array of write cursors.
for (int id = int(es.size()); id-- > 0;) {
const auto &[cst, u, to] = es[id];
if (u != to && to != rt) {
ids[--start[to]] = id;
}
}
// Parallel edges receive identical reductions in every contracted heap.
// Keep the cheapest one, breaking ties by the original input index.
// If there are n-1 eligible edges, the precheck guarantees unique rows.
if (ids.size() > std::size_t(n - 1)) {
std::vector<int> seen(n, -1);
std::vector<std::size_t> best(n);
std::size_t kept = 0;
for (int v = 0; v < n; ++v) {
auto l = start[v], r = start[v + 1];
start[v] = kept;
for (auto i = l; i < r; ++i) {
int id = ids[i], u = std::get<1>(es[id]);
if (seen[u] != v) {
seen[u] = v;
best[u] = kept;
ids[kept++] = id;
} else if (std::get<0>(es[id]) < std::get<0>(es[ids[best[u]]])) {
ids[best[u]] = id;
}
}
}
start[n] = kept;
ids.resize(kept);
}
buf.resize(ids.size());
std::vector<heap_node *> in(cap, nullptr), roots;
for (int v = 0; v < n; ++v) {
roots.clear();
for (auto i = start[v]; i < start[v + 1]; ++i) {
int id = ids[i];
const auto &[cst, u, to] = es[id];
roots.push_back(make_node(cst, u, id));
}
// Pair adjacent heaps in rounds: linear total construction work.
for (auto r = roots.size(); r > 1;) {
std::size_t next = 0;
for (std::size_t i = 0; i < r; i += 2) {
roots[next++] = i + 1 < r ? meld(roots[i], roots[i + 1]) : roots[i];
}
r = next;
}
if (!roots.empty()) {
in[v] = roots[0];
}
}
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);
st[rt] = 2;
int cnt = n;
std::vector<int> pth;
for (int s = 0; s < n; s++) {
if (st[s] != 0) {
continue;
}
pth.clear();
pth.push_back(s);
while (true) {
int com = pth.back();
st[com] = 1;
edge bst;
int pre = -1;
while (in[com] != nullptr) {
edge can = pop(in[com]);
int from = fd(can.u);
// Active path entries are already component representatives.
if (from != com) {
bst = can;
pre = from;
break;
}
}
if (pre == -1) {
return std::nullopt;
}
sel[com] = bst;
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[mem] = ctr;
cp[mem] = ctr;
st[mem] = 2;
if (mem == pre) {
break;
}
}
pth.push_back(ctr);
}
for (int com : pth) {
st[com] = 2;
}
}
std::vector<char> ext(cnt);
ext[rt] = true;
std::vector<int> eid;
eid.reserve(n - 1);
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 rt.
/// Edges are (cost, from, to); loops and edges entering rt are ignored.
/// Parallel edges keep the minimum (cost, original index).
/// Requires 0 < n <= INT_MAX/2, valid vertices, and at most INT_MAX edges.
/// T needs zero/default construction, copying, <, unary -, and += with
/// ordered additive semantics; all intermediate values must be representable.
/// @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
// NOYA_DIRECTED_MST_HPP
#ifndef NOYA_DIRECTED_MST_HPP
#define NOYA_DIRECTED_MST_HPP 1
/// @complexity Time: O(n + m + M log(n + 1)), M distinct eligible edge pairs.
/// Space: O(n + m).
#include <cassert>
#include <cstddef>
#include <limits>
#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;
bool dirty = false;
};
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_) {}
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;
nd->dirty = true;
}
return nd;
}
static void push(heap_node *nd) {
if (nd == nullptr || !nd->dirty) {
return;
}
add(nd->l, nd->lz);
add(nd->r, nd->lz);
nd->lz = T{};
nd->dirty = false;
}
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(0 < n && n <= std::numeric_limits<int>::max() / 2);
assert(0 <= rt && rt < n);
assert(es.size() <= std::size_t(std::numeric_limits<int>::max()));
const int cap = 2 * n;
std::vector<std::size_t> start(std::size_t(n) + 1);
for (const auto &[cst, u, to] : es) {
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to && to != rt) {
++start[to];
}
}
for (int v = 0; v < n; ++v) {
if (v != rt && start[v] == 0) {
return std::nullopt;
}
}
if (n == 1) {
return std::pair<T, std::vector<int>>{T{}, {}};
}
for (int v = 0; v < n; ++v) {
start[v + 1] += start[v];
}
std::vector<int> ids(start.back());
// Stable incoming-edge rows, without a second array of write cursors.
for (int id = int(es.size()); id-- > 0;) {
const auto &[cst, u, to] = es[id];
if (u != to && to != rt) {
ids[--start[to]] = id;
}
}
// Parallel edges receive identical reductions in every contracted heap.
// Keep the cheapest one, breaking ties by the original input index.
// If there are n-1 eligible edges, the precheck guarantees unique rows.
if (ids.size() > std::size_t(n - 1)) {
std::vector<int> seen(n, -1);
std::vector<std::size_t> best(n);
std::size_t kept = 0;
for (int v = 0; v < n; ++v) {
auto l = start[v], r = start[v + 1];
start[v] = kept;
for (auto i = l; i < r; ++i) {
int id = ids[i], u = std::get<1>(es[id]);
if (seen[u] != v) {
seen[u] = v;
best[u] = kept;
ids[kept++] = id;
} else if (std::get<0>(es[id]) < std::get<0>(es[ids[best[u]]])) {
ids[best[u]] = id;
}
}
}
start[n] = kept;
ids.resize(kept);
}
buf.resize(ids.size());
std::vector<heap_node *> in(cap, nullptr), roots;
for (int v = 0; v < n; ++v) {
roots.clear();
for (auto i = start[v]; i < start[v + 1]; ++i) {
int id = ids[i];
const auto &[cst, u, to] = es[id];
roots.push_back(make_node(cst, u, id));
}
// Pair adjacent heaps in rounds: linear total construction work.
for (auto r = roots.size(); r > 1;) {
std::size_t next = 0;
for (std::size_t i = 0; i < r; i += 2) {
roots[next++] = i + 1 < r ? meld(roots[i], roots[i + 1]) : roots[i];
}
r = next;
}
if (!roots.empty()) {
in[v] = roots[0];
}
}
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);
st[rt] = 2;
int cnt = n;
std::vector<int> pth;
for (int s = 0; s < n; s++) {
if (st[s] != 0) {
continue;
}
pth.clear();
pth.push_back(s);
while (true) {
int com = pth.back();
st[com] = 1;
edge bst;
int pre = -1;
while (in[com] != nullptr) {
edge can = pop(in[com]);
int from = fd(can.u);
// Active path entries are already component representatives.
if (from != com) {
bst = can;
pre = from;
break;
}
}
if (pre == -1) {
return std::nullopt;
}
sel[com] = bst;
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[mem] = ctr;
cp[mem] = ctr;
st[mem] = 2;
if (mem == pre) {
break;
}
}
pth.push_back(ctr);
}
for (int com : pth) {
st[com] = 2;
}
}
std::vector<char> ext(cnt);
ext[rt] = true;
std::vector<int> eid;
eid.reserve(n - 1);
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 rt.
/// Edges are (cost, from, to); loops and edges entering rt are ignored.
/// Parallel edges keep the minimum (cost, original index).
/// Requires 0 < n <= INT_MAX/2, valid vertices, and at most INT_MAX edges.
/// T needs zero/default construction, copying, <, unary -, and += with
/// ordered additive semantics; all intermediate values must be representable.
/// @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
/// @complexity Time: O(n + m + M log(n + 1)), M distinct eligible edge pairs.
/// Space: O(n + m).
#include <cassert>
#include <cstddef>
#include <limits>
#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;
bool dirty = false;
};
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_) {}
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;
nd->dirty = true;
}
return nd;
}
static void push(heap_node *nd) {
if (nd == nullptr || !nd->dirty) {
return;
}
add(nd->l, nd->lz);
add(nd->r, nd->lz);
nd->lz = T{};
nd->dirty = false;
}
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(0 < n && n <= std::numeric_limits<int>::max() / 2);
assert(0 <= rt && rt < n);
assert(es.size() <= std::size_t(std::numeric_limits<int>::max()));
const int cap = 2 * n;
std::vector<std::size_t> start(std::size_t(n) + 1);
for (const auto &[cst, u, to] : es) {
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to && to != rt) {
++start[to];
}
}
for (int v = 0; v < n; ++v) {
if (v != rt && start[v] == 0) {
return std::nullopt;
}
}
if (n == 1) {
return std::pair<T, std::vector<int>>{T{}, {}};
}
for (int v = 0; v < n; ++v) {
start[v + 1] += start[v];
}
std::vector<int> ids(start.back());
// Stable incoming-edge rows, without a second array of write cursors.
for (int id = int(es.size()); id-- > 0;) {
const auto &[cst, u, to] = es[id];
if (u != to && to != rt) {
ids[--start[to]] = id;
}
}
// Parallel edges receive identical reductions in every contracted heap.
// Keep the cheapest one, breaking ties by the original input index.
// If there are n-1 eligible edges, the precheck guarantees unique rows.
if (ids.size() > std::size_t(n - 1)) {
std::vector<int> seen(n, -1);
std::vector<std::size_t> best(n);
std::size_t kept = 0;
for (int v = 0; v < n; ++v) {
auto l = start[v], r = start[v + 1];
start[v] = kept;
for (auto i = l; i < r; ++i) {
int id = ids[i], u = std::get<1>(es[id]);
if (seen[u] != v) {
seen[u] = v;
best[u] = kept;
ids[kept++] = id;
} else if (std::get<0>(es[id]) < std::get<0>(es[ids[best[u]]])) {
ids[best[u]] = id;
}
}
}
start[n] = kept;
ids.resize(kept);
}
buf.resize(ids.size());
std::vector<heap_node *> in(cap, nullptr), roots;
for (int v = 0; v < n; ++v) {
roots.clear();
for (auto i = start[v]; i < start[v + 1]; ++i) {
int id = ids[i];
const auto &[cst, u, to] = es[id];
roots.push_back(make_node(cst, u, id));
}
// Pair adjacent heaps in rounds: linear total construction work.
for (auto r = roots.size(); r > 1;) {
std::size_t next = 0;
for (std::size_t i = 0; i < r; i += 2) {
roots[next++] = i + 1 < r ? meld(roots[i], roots[i + 1]) : roots[i];
}
r = next;
}
if (!roots.empty()) {
in[v] = roots[0];
}
}
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);
st[rt] = 2;
int cnt = n;
std::vector<int> pth;
for (int s = 0; s < n; s++) {
if (st[s] != 0) {
continue;
}
pth.clear();
pth.push_back(s);
while (true) {
int com = pth.back();
st[com] = 1;
edge bst;
int pre = -1;
while (in[com] != nullptr) {
edge can = pop(in[com]);
int from = fd(can.u);
// Active path entries are already component representatives.
if (from != com) {
bst = can;
pre = from;
break;
}
}
if (pre == -1) {
return std::nullopt;
}
sel[com] = bst;
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[mem] = ctr;
cp[mem] = ctr;
st[mem] = 2;
if (mem == pre) {
break;
}
}
pth.push_back(ctr);
}
for (int com : pth) {
st[com] = 2;
}
}
std::vector<char> ext(cnt);
ext[rt] = true;
std::vector<int> eid;
eid.reserve(n - 1);
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 rt.
/// Edges are (cost, from, to); loops and edges entering rt are ignored.
/// Parallel edges keep the minimum (cost, original index).
/// Requires 0 < n <= INT_MAX/2, valid vertices, and at most INT_MAX edges.
/// T needs zero/default construction, copying, <, unary -, and += with
/// ordered additive semantics; all intermediate values must be representable.
/// @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
// NOYA_DIRECTED_MST_HPP