directed_mst.hpp¶
Minimum spanning arborescence rooted at root in O(m log n). Edges are (cost, from, to).
Verified by directedmst.
求指定根可达所有点的最小权有向生成树（最小树形图），并恢复入边选择。
Implementation¶
#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 cost{};
int from = -1;
int id = -1;
};
struct heap_node {
heap_node *left = nullptr;
heap_node *right = nullptr;
edge value;
T lazy{};
int rank = 1;
};
int n;
int root;
const std::vector<std::tuple<T, int, int>> &edges;
std::vector<heap_node> pool;
int pool_index = 0;
solver(int n_, int root_, const std::vector<std::tuple<T, int, int>> &edges_)
: n(n_), root(root_), edges(edges_), pool(edges.size()) {}
static int heap_rank(heap_node *node) {
return node == nullptr ? 0 : node->rank;
}
static heap_node *add(heap_node *node, const T &delta) {
if (node != nullptr) {
node->value.cost += delta;
node->lazy += delta;
}
return node;
}
static void push(heap_node *node) {
if (node == nullptr) {
return;
}
add(node->left, node->lazy);
add(node->right, node->lazy);
node->lazy = T{};
}
static heap_node *meld(heap_node *first, heap_node *second) {
if (first == nullptr) {
return second;
}
if (second == nullptr) {
return first;
}
if (std::pair(second->value.cost, second->value.id) <
std::pair(first->value.cost, first->value.id)) {
std::swap(first, second);
}
push(first);
first->right = meld(first->right, second);
if (heap_rank(first->left) < heap_rank(first->right)) {
std::swap(first->left, first->right);
}
first->rank = heap_rank(first->right) + 1;
return first;
}
heap_node *make_node(T cost, int from, int id) {
heap_node *node = &pool[pool_index++];
node->value = {cost, from, id};
return node;
}
static edge pop(heap_node *&node) {
push(node);
edge result = node->value;
node = meld(node->left, node->right);
return result;
}
std::optional<std::pair<T, std::vector<int>>> run() {
assert(n > 0);
assert(0 <= root && root < n);
const int max_components = 2 * n;
std::vector<heap_node *> incoming(max_components, nullptr);
for (int id = 0; id < int(edges.size()); id++) {
auto [cost, from, to] = edges[id];
assert(0 <= from && from < n);
assert(0 <= to && to < n);
if (from != to) {
incoming[to] = meld(incoming[to], make_node(cost, from, id));
}
}
std::vector<int> dsu(max_components);
std::iota(dsu.begin(), dsu.end(), 0);
auto find = [&](int vertex) {
int representative = vertex;
while (dsu[representative] != representative) {
representative = dsu[representative];
}
while (dsu[vertex] != vertex) {
int next = dsu[vertex];
dsu[vertex] = representative;
vertex = next;
}
return representative;
};
std::vector<char> state(max_components);
std::vector<edge> chosen(max_components);
std::vector<int> contraction_parent(max_components, -1);
std::vector<int> representative(max_components);
std::iota(representative.begin(), representative.end(), 0);
state[root] = 2;
int component_count = n;
for (int start = 0; start < n; start++) {
if (state[start] != 0) {
continue;
}
std::vector<int> path = {start};
while (true) {
int component = path.back();
state[component] = 1;
edge best;
bool found = false;
while (incoming[component] != nullptr) {
edge candidate = pop(incoming[component]);
if (find(candidate.from) != find(component)) {
best = candidate;
found = true;
break;
}
}
if (!found) {
return std::nullopt;
}
chosen[component] = best;
int previous = representative[find(best.from)];
if (state[previous] == 0) {
path.push_back(previous);
continue;
}
if (state[previous] == 2) {
break;
}
int contracted = component_count++;
assert(contracted < max_components);
while (true) {
int member = path.back();
path.pop_back();
incoming[contracted] =
meld(incoming[contracted],
add(incoming[member], -chosen[member].cost));
dsu[find(member)] = contracted;
contraction_parent[member] = contracted;
state[member] = 2;
if (member == previous) {
break;
}
}
representative[find(contracted)] = contracted;
path.push_back(contracted);
}
for (int component : path) {
state[component] = 2;
}
}
std::vector<char> expanded(component_count);
expanded[root] = true;
std::vector<int> edge_ids;
for (int component = component_count - 1; component >= 0; component--) {
if (expanded[component]) {
continue;
}
int id = chosen[component].id;
if (id == -1) {
return std::nullopt;
}
edge_ids.push_back(id);
int vertex = std::get<2>(edges[id]);
while (vertex != -1 && !expanded[vertex]) {
expanded[vertex] = true;
vertex = contraction_parent[vertex];
}
}
if (int(edge_ids.size()) != n - 1) {
return std::nullopt;
}
T cost{};
for (int id : edge_ids) {
cost += std::get<0>(edges[id]);
}
return std::pair<T, std::vector<int>>{cost, std::move(edge_ids)};
}
};
} // 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 root,
const std::vector<std::tuple<T, int, int>> &edges) {
return directed_mst_internal::solver<T>(n, root, edges).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 cost{};
int from = -1;
int id = -1;
};
struct heap_node {
heap_node *left = nullptr;
heap_node *right = nullptr;
edge value;
T lazy{};
int rank = 1;
};
int n;
int root;
const std::vector<std::tuple<T, int, int>> &edges;
std::vector<heap_node> pool;
int pool_index = 0;
solver(int n_, int root_, const std::vector<std::tuple<T, int, int>> &edges_)
: n(n_), root(root_), edges(edges_), pool(edges.size()) {}
static int heap_rank(heap_node *node) {
return node == nullptr ? 0 : node->rank;
}
static heap_node *add(heap_node *node, const T &delta) {
if (node != nullptr) {
node->value.cost += delta;
node->lazy += delta;
}
return node;
}
static void push(heap_node *node) {
if (node == nullptr) {
return;
}
add(node->left, node->lazy);
add(node->right, node->lazy);
node->lazy = T{};
}
static heap_node *meld(heap_node *first, heap_node *second) {
if (first == nullptr) {
return second;
}
if (second == nullptr) {
return first;
}
if (std::pair(second->value.cost, second->value.id) <
std::pair(first->value.cost, first->value.id)) {
std::swap(first, second);
}
push(first);
first->right = meld(first->right, second);
if (heap_rank(first->left) < heap_rank(first->right)) {
std::swap(first->left, first->right);
}
first->rank = heap_rank(first->right) + 1;
return first;
}
heap_node *make_node(T cost, int from, int id) {
heap_node *node = &pool[pool_index++];
node->value = {cost, from, id};
return node;
}
static edge pop(heap_node *&node) {
push(node);
edge result = node->value;
node = meld(node->left, node->right);
return result;
}
std::optional<std::pair<T, std::vector<int>>> run() {
assert(n > 0);
assert(0 <= root && root < n);
const int max_components = 2 * n;
std::vector<heap_node *> incoming(max_components, nullptr);
for (int id = 0; id < int(edges.size()); id++) {
auto [cost, from, to] = edges[id];
assert(0 <= from && from < n);
assert(0 <= to && to < n);
if (from != to) {
incoming[to] = meld(incoming[to], make_node(cost, from, id));
}
}
std::vector<int> dsu(max_components);
std::iota(dsu.begin(), dsu.end(), 0);
auto find = [&](int vertex) {
int representative = vertex;
while (dsu[representative] != representative) {
representative = dsu[representative];
}
while (dsu[vertex] != vertex) {
int next = dsu[vertex];
dsu[vertex] = representative;
vertex = next;
}
return representative;
};
std::vector<char> state(max_components);
std::vector<edge> chosen(max_components);
std::vector<int> contraction_parent(max_components, -1);
std::vector<int> representative(max_components);
std::iota(representative.begin(), representative.end(), 0);
state[root] = 2;
int component_count = n;
for (int start = 0; start < n; start++) {
if (state[start] != 0) {
continue;
}
std::vector<int> path = {start};
while (true) {
int component = path.back();
state[component] = 1;
edge best;
bool found = false;
while (incoming[component] != nullptr) {
edge candidate = pop(incoming[component]);
if (find(candidate.from) != find(component)) {
best = candidate;
found = true;
break;
}
}
if (!found) {
return std::nullopt;
}
chosen[component] = best;
int previous = representative[find(best.from)];
if (state[previous] == 0) {
path.push_back(previous);
continue;
}
if (state[previous] == 2) {
break;
}
int contracted = component_count++;
assert(contracted < max_components);
while (true) {
int member = path.back();
path.pop_back();
incoming[contracted] =
meld(incoming[contracted],
add(incoming[member], -chosen[member].cost));
dsu[find(member)] = contracted;
contraction_parent[member] = contracted;
state[member] = 2;
if (member == previous) {
break;
}
}
representative[find(contracted)] = contracted;
path.push_back(contracted);
}
for (int component : path) {
state[component] = 2;
}
}
std::vector<char> expanded(component_count);
expanded[root] = true;
std::vector<int> edge_ids;
for (int component = component_count - 1; component >= 0; component--) {
if (expanded[component]) {
continue;
}
int id = chosen[component].id;
if (id == -1) {
return std::nullopt;
}
edge_ids.push_back(id);
int vertex = std::get<2>(edges[id]);
while (vertex != -1 && !expanded[vertex]) {
expanded[vertex] = true;
vertex = contraction_parent[vertex];
}
}
if (int(edge_ids.size()) != n - 1) {
return std::nullopt;
}
T cost{};
for (int id : edge_ids) {
cost += std::get<0>(edges[id]);
}
return std::pair<T, std::vector<int>>{cost, std::move(edge_ids)};
}
};
} // 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 root,
const std::vector<std::tuple<T, int, int>> &edges) {
return directed_mst_internal::solver<T>(n, root, edges).run();
}
} // namespace noya