minimum_spanning_tree.hpp¶
求最小/最大生成森林或 dense Prim 森林，边按 (权值,起点,终点) 传入；返回 int64 总权与原边编号，内部使用 128 位累加。
Complexity: Time: \(O(V + E \log (E + 1))\) Kruskal, \(O(V + E \alpha (V + 1))\) for ordered edges; \(O(V^{2} + E)\) dense Prim. Space: \(O(V + E)\) Kruskal; \(O(V^{2} + V)\) dense Prim.
AC 记录：minimum_spanning_tree。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(V+E log(E+1)) Kruskal, O(V+E alpha(V+1)) for ordered edges;
/// O(V^2+E) dense Prim. Space: O(V+E) Kruskal; O(V^2+V) dense Prim.
#include "atcoder/dsu.hpp"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <functional>
#include <limits>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
namespace internal {
inline int64_t spanning_forest_cost(__int128 cost) {
assert(cost >= std::numeric_limits<int64_t>::min() &&
cost <= std::numeric_limits<int64_t>::max());
return int64_t(cost);
}
template <class T, class Compare>
std::pair<int64_t, std::vector<int>>
kruskal_spanning_forest(int n, const std::vector<T> &es, Compare order) {
assert(n >= 0 && es.size() <= std::size_t(std::numeric_limits<int>::max()));
auto key = [](const T &edge) {
const auto &[w, u, v] = edge;
return std::tuple{w, int(u), int(v)};
};
using Key = decltype(key(es[0]));
struct Record { Key key; int id; };
std::vector<Record> p;
p.reserve(es.size());
for (int i = 0; i < int(es.size()); ++i) {
const auto &[w, u, v] = es[i];
assert(0 <= u && u < n && 0 <= v && v < n);
if (u != v) p.push_back({key(es[i]), i});
}
if (n <= 1 || p.empty()) return {0, {}};
auto cmp = [&](const auto &a, const auto &b) {
return order(std::get<0>(a.key), std::get<0>(b.key));
};
if (!std::is_sorted(p.begin(), p.end(), cmp))
std::sort(p.begin(), p.end(), cmp);
atcoder::dsu dsu(n);
__int128 cost = 0;
std::vector<int> selected;
selected.reserve(std::min(std::size_t(n - 1), p.size()));
for (const auto &[edge, i] : p) {
const auto &[w, u, v] = edge;
int a = dsu.leader(u), b = dsu.leader(v);
if (a == b) continue;
dsu.merge(a, b);
cost += static_cast<__int128>(w);
selected.push_back(i);
if (int(selected.size()) == n - 1) break;
}
return {spanning_forest_cost(cost), std::move(selected)};
}
} // namespace internal
/// @brief Minimum spanning forest by Kruskal. Edges are (weight,u,v).
/// Return (int64 total, original edge IDs), without modifying the input.
/// Integral weights and intermediate sums must fit signed __int128, and the
/// final total must fit int64_t. Loops are ignored; equal optima may differ.
template <class T>
std::pair<int64_t, std::vector<int>>
minimum_spanning_tree(int n, const std::vector<T> &es) {
return internal::kruskal_spanning_forest(n, es, std::less<>{});
}
/// @brief Maximum spanning forest, with the same contract as the minimum.
template <class T>
std::pair<int64_t, std::vector<int>>
maximum_spanning_tree(int n, const std::vector<T> &es) {
return internal::kruskal_spanning_forest(n, es, std::greater<>{});
}
/// @brief Minimum spanning forest by dense Prim; restart in each component.
/// Uses the same edge layout, numeric and ownership contracts as Kruskal.
template <class T>
std::pair<int64_t, std::vector<int>> prim_dense(int n, const std::vector<T> &es) {
assert(n >= 0 && es.size() <= std::size_t(std::numeric_limits<int>::max()));
if (es.empty()) return {0, {}};
std::vector<int> id(std::size_t(n) * n, -1);
int nonloops = 0;
for (int i = 0; i < int(es.size()); ++i) {
const auto &[w, u, v] = es[i];
assert(0 <= u && u < n && 0 <= v && v < n);
if (u == v) continue;
++nonloops;
auto &a = id[std::size_t(u) * n + v];
if (a != -1) {
const auto &[old_weight, old_u, old_v] = es[a];
if (!(std::tie(w, u, v) < std::tie(old_weight, old_u, old_v))) continue;
}
a = id[std::size_t(v) * n + u] = i;
}
if (!nonloops) return {0, {}};
std::vector<int> selected, best(n, -1);
selected.reserve(std::min(n - 1, nonloops));
std::vector<unsigned char> used(n);
auto better = [&](int a, int b) {
if (a == -1) return b;
if (b == -1) return a;
const auto &[x, u1, v1] = es[a];
const auto &[y, u2, v2] = es[b];
return x < y ? a : b;
};
__int128 cost = 0;
for (int root = 0; root < n; ++root) if (!used[root]) {
int u = root;
while (true) {
used[u] = true;
int next = -1;
const int *row = id.data() + std::size_t(u) * n;
for (int v = 0; v < n; ++v) if (!used[v]) {
best[v] = better(best[v], row[v]);
next = better(next, best[v]);
}
if (next == -1) break;
selected.push_back(next);
const auto &[w, a, b] = es[next];
cost += static_cast<__int128>(w);
u = used[a] ? b : a;
}
}
return {internal::spanning_forest_cost(cost), std::move(selected)};
}
} // namespace noya
// NOYA_MINIMUM_SPANNING_TREE_HPP
#ifndef NOYA_MINIMUM_SPANNING_TREE_HPP
#define NOYA_MINIMUM_SPANNING_TREE_HPP 1
/// @complexity Time: O(V+E log(E+1)) Kruskal, O(V+E alpha(V+1)) for ordered edges;
/// O(V^2+E) dense Prim. Space: O(V+E) Kruskal; O(V^2+V) dense Prim.
#include "atcoder/dsu.hpp"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <functional>
#include <limits>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
namespace internal {
inline int64_t spanning_forest_cost(__int128 cost) {
assert(cost >= std::numeric_limits<int64_t>::min() &&
cost <= std::numeric_limits<int64_t>::max());
return int64_t(cost);
}
template <class T, class Compare>
std::pair<int64_t, std::vector<int>>
kruskal_spanning_forest(int n, const std::vector<T> &es, Compare order) {
assert(n >= 0 && es.size() <= std::size_t(std::numeric_limits<int>::max()));
auto key = [](const T &edge) {
const auto &[w, u, v] = edge;
return std::tuple{w, int(u), int(v)};
};
using Key = decltype(key(es[0]));
struct Record { Key key; int id; };
std::vector<Record> p;
p.reserve(es.size());
for (int i = 0; i < int(es.size()); ++i) {
const auto &[w, u, v] = es[i];
assert(0 <= u && u < n && 0 <= v && v < n);
if (u != v) p.push_back({key(es[i]), i});
}
if (n <= 1 || p.empty()) return {0, {}};
auto cmp = [&](const auto &a, const auto &b) {
return order(std::get<0>(a.key), std::get<0>(b.key));
};
if (!std::is_sorted(p.begin(), p.end(), cmp))
std::sort(p.begin(), p.end(), cmp);
atcoder::dsu dsu(n);
__int128 cost = 0;
std::vector<int> selected;
selected.reserve(std::min(std::size_t(n - 1), p.size()));
for (const auto &[edge, i] : p) {
const auto &[w, u, v] = edge;
int a = dsu.leader(u), b = dsu.leader(v);
if (a == b) continue;
dsu.merge(a, b);
cost += static_cast<__int128>(w);
selected.push_back(i);
if (int(selected.size()) == n - 1) break;
}
return {spanning_forest_cost(cost), std::move(selected)};
}
} // namespace internal
/// @brief Minimum spanning forest by Kruskal. Edges are (weight,u,v).
/// Return (int64 total, original edge IDs), without modifying the input.
/// Integral weights and intermediate sums must fit signed __int128, and the
/// final total must fit int64_t. Loops are ignored; equal optima may differ.
template <class T>
std::pair<int64_t, std::vector<int>>
minimum_spanning_tree(int n, const std::vector<T> &es) {
return internal::kruskal_spanning_forest(n, es, std::less<>{});
}
/// @brief Maximum spanning forest, with the same contract as the minimum.
template <class T>
std::pair<int64_t, std::vector<int>>
maximum_spanning_tree(int n, const std::vector<T> &es) {
return internal::kruskal_spanning_forest(n, es, std::greater<>{});
}
/// @brief Minimum spanning forest by dense Prim; restart in each component.
/// Uses the same edge layout, numeric and ownership contracts as Kruskal.
template <class T>
std::pair<int64_t, std::vector<int>> prim_dense(int n, const std::vector<T> &es) {
assert(n >= 0 && es.size() <= std::size_t(std::numeric_limits<int>::max()));
if (es.empty()) return {0, {}};
std::vector<int> id(std::size_t(n) * n, -1);
int nonloops = 0;
for (int i = 0; i < int(es.size()); ++i) {
const auto &[w, u, v] = es[i];
assert(0 <= u && u < n && 0 <= v && v < n);
if (u == v) continue;
++nonloops;
auto &a = id[std::size_t(u) * n + v];
if (a != -1) {
const auto &[old_weight, old_u, old_v] = es[a];
if (!(std::tie(w, u, v) < std::tie(old_weight, old_u, old_v))) continue;
}
a = id[std::size_t(v) * n + u] = i;
}
if (!nonloops) return {0, {}};
std::vector<int> selected, best(n, -1);
selected.reserve(std::min(n - 1, nonloops));
std::vector<unsigned char> used(n);
auto better = [&](int a, int b) {
if (a == -1) return b;
if (b == -1) return a;
const auto &[x, u1, v1] = es[a];
const auto &[y, u2, v2] = es[b];
return x < y ? a : b;
};
__int128 cost = 0;
for (int root = 0; root < n; ++root) if (!used[root]) {
int u = root;
while (true) {
used[u] = true;
int next = -1;
const int *row = id.data() + std::size_t(u) * n;
for (int v = 0; v < n; ++v) if (!used[v]) {
best[v] = better(best[v], row[v]);
next = better(next, best[v]);
}
if (next == -1) break;
selected.push_back(next);
const auto &[w, a, b] = es[next];
cost += static_cast<__int128>(w);
u = used[a] ? b : a;
}
}
return {internal::spanning_forest_cost(cost), std::move(selected)};
}
} // namespace noya
#endif // NOYA_MINIMUM_SPANNING_TREE_HPP
/// @complexity Time: O(V+E log(E+1)) Kruskal, O(V+E alpha(V+1)) for ordered edges;
/// O(V^2+E) dense Prim. Space: O(V+E) Kruskal; O(V^2+V) dense Prim.
#include <algorithm>
#include <cassert>
#include <vector>
namespace atcoder {
// Implement (union by size) + (path compression)
// Reference:
// Zvi Galil and Giuseppe F. Italiano,
// Data structures and algorithms for disjoint set union problems
struct dsu {
public:
dsu() : _n(0) {}
explicit dsu(int n) : _n(n), parent_or_size(n, -1) {}
int merge(int a, int b) {
assert(0 <= a && a < _n);
assert(0 <= b && b < _n);
int x = leader(a), y = leader(b);
if (x == y) return x;
if (-parent_or_size[x] < -parent_or_size[y]) std::swap(x, y);
parent_or_size[x] += parent_or_size[y];
parent_or_size[y] = x;
return x;
}
bool same(int a, int b) {
assert(0 <= a && a < _n);
assert(0 <= b && b < _n);
return leader(a) == leader(b);
}
int leader(int a) {
assert(0 <= a && a < _n);
return _leader(a);
}
int size(int a) {
assert(0 <= a && a < _n);
return -parent_or_size[leader(a)];
}
std::vector<std::vector<int>> groups() {
std::vector<int> leader_buf(_n), group_size(_n);
for (int i = 0; i < _n; i++) {
leader_buf[i] = leader(i);
group_size[leader_buf[i]]++;
}
std::vector<std::vector<int>> result(_n);
for (int i = 0; i < _n; i++) {
result[i].reserve(group_size[i]);
}
for (int i = 0; i < _n; i++) {
result[leader_buf[i]].push_back(i);
}
result.erase(
std::remove_if(result.begin(), result.end(),
[&](const std::vector<int>& v) { return v.empty(); }),
result.end());
return result;
}
private:
int _n;
// root node: -1 * component size
// otherwise: parent
std::vector<int> parent_or_size;
int _leader(int a) {
if (parent_or_size[a] < 0) return a;
return parent_or_size[a] = _leader(parent_or_size[a]);
}
};
} // namespace atcoder
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <functional>
#include <limits>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
namespace internal {
inline int64_t spanning_forest_cost(__int128 cost) {
assert(cost >= std::numeric_limits<int64_t>::min() &&
cost <= std::numeric_limits<int64_t>::max());
return int64_t(cost);
}
template <class T, class Compare>
std::pair<int64_t, std::vector<int>>
kruskal_spanning_forest(int n, const std::vector<T> &es, Compare order) {
assert(n >= 0 && es.size() <= std::size_t(std::numeric_limits<int>::max()));
auto key = [](const T &edge) {
const auto &[w, u, v] = edge;
return std::tuple{w, int(u), int(v)};
};
using Key = decltype(key(es[0]));
struct Record { Key key; int id; };
std::vector<Record> p;
p.reserve(es.size());
for (int i = 0; i < int(es.size()); ++i) {
const auto &[w, u, v] = es[i];
assert(0 <= u && u < n && 0 <= v && v < n);
if (u != v) p.push_back({key(es[i]), i});
}
if (n <= 1 || p.empty()) return {0, {}};
auto cmp = [&](const auto &a, const auto &b) {
return order(std::get<0>(a.key), std::get<0>(b.key));
};
if (!std::is_sorted(p.begin(), p.end(), cmp))
std::sort(p.begin(), p.end(), cmp);
atcoder::dsu dsu(n);
__int128 cost = 0;
std::vector<int> selected;
selected.reserve(std::min(std::size_t(n - 1), p.size()));
for (const auto &[edge, i] : p) {
const auto &[w, u, v] = edge;
int a = dsu.leader(u), b = dsu.leader(v);
if (a == b) continue;
dsu.merge(a, b);
cost += static_cast<__int128>(w);
selected.push_back(i);
if (int(selected.size()) == n - 1) break;
}
return {spanning_forest_cost(cost), std::move(selected)};
}
} // namespace internal
/// @brief Minimum spanning forest by Kruskal. Edges are (weight,u,v).
/// Return (int64 total, original edge IDs), without modifying the input.
/// Integral weights and intermediate sums must fit signed __int128, and the
/// final total must fit int64_t. Loops are ignored; equal optima may differ.
template <class T>
std::pair<int64_t, std::vector<int>>
minimum_spanning_tree(int n, const std::vector<T> &es) {
return internal::kruskal_spanning_forest(n, es, std::less<>{});
}
/// @brief Maximum spanning forest, with the same contract as the minimum.
template <class T>
std::pair<int64_t, std::vector<int>>
maximum_spanning_tree(int n, const std::vector<T> &es) {
return internal::kruskal_spanning_forest(n, es, std::greater<>{});
}
/// @brief Minimum spanning forest by dense Prim; restart in each component.
/// Uses the same edge layout, numeric and ownership contracts as Kruskal.
template <class T>
std::pair<int64_t, std::vector<int>> prim_dense(int n, const std::vector<T> &es) {
assert(n >= 0 && es.size() <= std::size_t(std::numeric_limits<int>::max()));
if (es.empty()) return {0, {}};
std::vector<int> id(std::size_t(n) * n, -1);
int nonloops = 0;
for (int i = 0; i < int(es.size()); ++i) {
const auto &[w, u, v] = es[i];
assert(0 <= u && u < n && 0 <= v && v < n);
if (u == v) continue;
++nonloops;
auto &a = id[std::size_t(u) * n + v];
if (a != -1) {
const auto &[old_weight, old_u, old_v] = es[a];
if (!(std::tie(w, u, v) < std::tie(old_weight, old_u, old_v))) continue;
}
a = id[std::size_t(v) * n + u] = i;
}
if (!nonloops) return {0, {}};
std::vector<int> selected, best(n, -1);
selected.reserve(std::min(n - 1, nonloops));
std::vector<unsigned char> used(n);
auto better = [&](int a, int b) {
if (a == -1) return b;
if (b == -1) return a;
const auto &[x, u1, v1] = es[a];
const auto &[y, u2, v2] = es[b];
return x < y ? a : b;
};
__int128 cost = 0;
for (int root = 0; root < n; ++root) if (!used[root]) {
int u = root;
while (true) {
used[u] = true;
int next = -1;
const int *row = id.data() + std::size_t(u) * n;
for (int v = 0; v < n; ++v) if (!used[v]) {
best[v] = better(best[v], row[v]);
next = better(next, best[v]);
}
if (next == -1) break;
selected.push_back(next);
const auto &[w, a, b] = es[next];
cost += static_cast<__int128>(w);
u = used[a] ? b : a;
}
}
return {internal::spanning_forest_cost(cost), std::move(selected)};
}
} // namespace noya
// NOYA_MINIMUM_SPANNING_TREE_HPP