special_spanning_trees.hpp¶
计算带特殊边类别或选择限制的生成树数量；适合需要在矩阵树定理上加入组合约束的题目。
Complexity: Time: Problem-dependent: parametric product tree uses repeated MSTs; Chebyshev MST is O(n^2 log n). Space: O(V + E), or O(n^2) for dense Chebyshev candidates.
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: Problem-dependent: parametric product tree uses repeated MSTs; Chebyshev MST is O(n^2 log n).
/// Space: O(V + E), or O(n^2) for dense Chebyshev candidates.
#include "atcoder/dsu.hpp"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <limits>
#include <numeric>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
struct bicriteria_spanning_tree_result {
std::int64_t s1 = 0;
std::int64_t s2 = 0;
std::vector<int> eid;
};
/// @brief Find a spanning tree minimizing (sum wa)*(sum
/// wb) by recursive supported-point search on the bicriteria
/// spanning-tree hull.
inline bicriteria_spanning_tree_result minimum_product_spanning_tree(
int n,
const std::vector<std::tuple<std::int64_t, std::int64_t, int, int>> &es) {
assert(n >= 1);
auto tr = [&](std::int64_t ca, std::int64_t cb) {
std::vector<int> ord(es.size());
std::iota(ord.begin(), ord.end(), 0);
std::sort(ord.begin(), ord.end(), [&](int a, int b) {
auto [a1, b1, u1, v1] = es[a];
auto [a2, b2, u2, v2] = es[b];
(void)u1;
(void)v1;
(void)u2;
(void)v2;
__int128 v11 = __int128(ca) * a1 + __int128(cb) * b1;
__int128 v21 = __int128(ca) * a2 + __int128(cb) * b2;
return v11 != v21 ? v11 < v21 : a < b;
});
atcoder::dsu dsu(n);
bicriteria_spanning_tree_result res;
for (int id : ord) {
auto [wa, wb, a, b] = es[id];
assert(0 <= a && a < n);
assert(0 <= b && b < n);
if (!dsu.same(a, b)) {
dsu.merge(a, b);
res.s1 += wa;
res.s2 += wb;
res.eid.push_back(id);
}
}
assert(int(res.eid.size()) == n - 1);
return res;
};
auto x = tr(1, 0);
auto y = tr(0, 1);
auto prd = [](const bicriteria_spanning_tree_result &val) {
return __int128(val.s1) * val.s2;
};
auto bst = prd(x) < prd(y) ? x : y;
auto dfs = [&](auto &self, const bicriteria_spanning_tree_result &l,
const bicriteria_spanning_tree_result &r) -> void {
std::int64_t ca = l.s2 - r.s2;
std::int64_t cb = r.s1 - l.s1;
auto mid = tr(ca, cb);
__int128 lv = __int128(ca) * l.s1 + __int128(cb) * l.s2;
__int128 mv = __int128(ca) * mid.s1 + __int128(cb) * mid.s2;
if (mv >= lv) {
return;
}
if (prd(mid) < prd(bst)) {
bst = mid;
}
self(self, l, mid);
self(self, mid, r);
};
dfs(dfs, x, y);
return bst;
}
/// @brief Exact O(n^2 log n) Chebyshev-distance MST of planar integer points,
/// returning total cost and complete-graph edge endpoint pairs.
inline std::pair<std::int64_t, std::vector<std::pair<int, int>>>
chebyshev_mst(const std::vector<std::pair<std::int64_t, std::int64_t>> &pt) {
int n = int(pt.size());
std::vector<std::tuple<std::int64_t, int, int>> es;
for (int a = 0; a < n; a++) {
for (int b = a + 1; b < n; b++) {
std::int64_t dis = std::max(std::abs(pt[a].first - pt[b].first),
std::abs(pt[a].second - pt[b].second));
es.emplace_back(dis, a, b);
}
}
std::sort(es.begin(), es.end());
atcoder::dsu dsu(n);
std::int64_t cst = 0;
std::vector<std::pair<int, int>> res;
for (auto [w, a, b] : es) {
if (!dsu.same(a, b)) {
dsu.merge(a, b);
cst += w;
res.emplace_back(a, b);
}
}
return {cst, res};
}
} // namespace noya
#ifndef NOYA_SPECIAL_SPANNING_TREES_HPP
#define NOYA_SPECIAL_SPANNING_TREES_HPP 1
/// @complexity Time: Problem-dependent: parametric product tree uses repeated MSTs; Chebyshev MST is O(n^2 log n).
/// Space: O(V + E), or O(n^2) for dense Chebyshev candidates.
#include "atcoder/dsu.hpp"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <limits>
#include <numeric>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
struct bicriteria_spanning_tree_result {
std::int64_t s1 = 0;
std::int64_t s2 = 0;
std::vector<int> eid;
};
/// @brief Find a spanning tree minimizing (sum wa)*(sum
/// wb) by recursive supported-point search on the bicriteria
/// spanning-tree hull.
inline bicriteria_spanning_tree_result minimum_product_spanning_tree(
int n,
const std::vector<std::tuple<std::int64_t, std::int64_t, int, int>> &es) {
assert(n >= 1);
auto tr = [&](std::int64_t ca, std::int64_t cb) {
std::vector<int> ord(es.size());
std::iota(ord.begin(), ord.end(), 0);
std::sort(ord.begin(), ord.end(), [&](int a, int b) {
auto [a1, b1, u1, v1] = es[a];
auto [a2, b2, u2, v2] = es[b];
(void)u1;
(void)v1;
(void)u2;
(void)v2;
__int128 v11 = __int128(ca) * a1 + __int128(cb) * b1;
__int128 v21 = __int128(ca) * a2 + __int128(cb) * b2;
return v11 != v21 ? v11 < v21 : a < b;
});
atcoder::dsu dsu(n);
bicriteria_spanning_tree_result res;
for (int id : ord) {
auto [wa, wb, a, b] = es[id];
assert(0 <= a && a < n);
assert(0 <= b && b < n);
if (!dsu.same(a, b)) {
dsu.merge(a, b);
res.s1 += wa;
res.s2 += wb;
res.eid.push_back(id);
}
}
assert(int(res.eid.size()) == n - 1);
return res;
};
auto x = tr(1, 0);
auto y = tr(0, 1);
auto prd = [](const bicriteria_spanning_tree_result &val) {
return __int128(val.s1) * val.s2;
};
auto bst = prd(x) < prd(y) ? x : y;
auto dfs = [&](auto &self, const bicriteria_spanning_tree_result &l,
const bicriteria_spanning_tree_result &r) -> void {
std::int64_t ca = l.s2 - r.s2;
std::int64_t cb = r.s1 - l.s1;
auto mid = tr(ca, cb);
__int128 lv = __int128(ca) * l.s1 + __int128(cb) * l.s2;
__int128 mv = __int128(ca) * mid.s1 + __int128(cb) * mid.s2;
if (mv >= lv) {
return;
}
if (prd(mid) < prd(bst)) {
bst = mid;
}
self(self, l, mid);
self(self, mid, r);
};
dfs(dfs, x, y);
return bst;
}
/// @brief Exact O(n^2 log n) Chebyshev-distance MST of planar integer points,
/// returning total cost and complete-graph edge endpoint pairs.
inline std::pair<std::int64_t, std::vector<std::pair<int, int>>>
chebyshev_mst(const std::vector<std::pair<std::int64_t, std::int64_t>> &pt) {
int n = int(pt.size());
std::vector<std::tuple<std::int64_t, int, int>> es;
for (int a = 0; a < n; a++) {
for (int b = a + 1; b < n; b++) {
std::int64_t dis = std::max(std::abs(pt[a].first - pt[b].first),
std::abs(pt[a].second - pt[b].second));
es.emplace_back(dis, a, b);
}
}
std::sort(es.begin(), es.end());
atcoder::dsu dsu(n);
std::int64_t cst = 0;
std::vector<std::pair<int, int>> res;
for (auto [w, a, b] : es) {
if (!dsu.same(a, b)) {
dsu.merge(a, b);
cst += w;
res.emplace_back(a, b);
}
}
return {cst, res};
}
} // namespace noya
#endif // NOYA_SPECIAL_SPANNING_TREES_HPP
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <limits>
#include <numeric>
#include <tuple>
#include <utility>
#include <vector>
/// @complexity Time: Problem-dependent: parametric product tree uses repeated MSTs; Chebyshev MST is O(n^2 log n).
/// Space: O(V + E), or O(n^2) for dense Chebyshev candidates.
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
namespace noya {
struct bicriteria_spanning_tree_result {
std::int64_t s1 = 0;
std::int64_t s2 = 0;
std::vector<int> eid;
};
/// @brief Find a spanning tree minimizing (sum wa)*(sum
/// wb) by recursive supported-point search on the bicriteria
/// spanning-tree hull.
inline bicriteria_spanning_tree_result minimum_product_spanning_tree(
int n,
const std::vector<std::tuple<std::int64_t, std::int64_t, int, int>> &es) {
assert(n >= 1);
auto tr = [&](std::int64_t ca, std::int64_t cb) {
std::vector<int> ord(es.size());
std::iota(ord.begin(), ord.end(), 0);
std::sort(ord.begin(), ord.end(), [&](int a, int b) {
auto [a1, b1, u1, v1] = es[a];
auto [a2, b2, u2, v2] = es[b];
(void)u1;
(void)v1;
(void)u2;
(void)v2;
__int128 v11 = __int128(ca) * a1 + __int128(cb) * b1;
__int128 v21 = __int128(ca) * a2 + __int128(cb) * b2;
return v11 != v21 ? v11 < v21 : a < b;
});
atcoder::dsu dsu(n);
bicriteria_spanning_tree_result res;
for (int id : ord) {
auto [wa, wb, a, b] = es[id];
assert(0 <= a && a < n);
assert(0 <= b && b < n);
if (!dsu.same(a, b)) {
dsu.merge(a, b);
res.s1 += wa;
res.s2 += wb;
res.eid.push_back(id);
}
}
assert(int(res.eid.size()) == n - 1);
return res;
};
auto x = tr(1, 0);
auto y = tr(0, 1);
auto prd = [](const bicriteria_spanning_tree_result &val) {
return __int128(val.s1) * val.s2;
};
auto bst = prd(x) < prd(y) ? x : y;
auto dfs = [&](auto &self, const bicriteria_spanning_tree_result &l,
const bicriteria_spanning_tree_result &r) -> void {
std::int64_t ca = l.s2 - r.s2;
std::int64_t cb = r.s1 - l.s1;
auto mid = tr(ca, cb);
__int128 lv = __int128(ca) * l.s1 + __int128(cb) * l.s2;
__int128 mv = __int128(ca) * mid.s1 + __int128(cb) * mid.s2;
if (mv >= lv) {
return;
}
if (prd(mid) < prd(bst)) {
bst = mid;
}
self(self, l, mid);
self(self, mid, r);
};
dfs(dfs, x, y);
return bst;
}
/// @brief Exact O(n^2 log n) Chebyshev-distance MST of planar integer points,
/// returning total cost and complete-graph edge endpoint pairs.
inline std::pair<std::int64_t, std::vector<std::pair<int, int>>>
chebyshev_mst(const std::vector<std::pair<std::int64_t, std::int64_t>> &pt) {
int n = int(pt.size());
std::vector<std::tuple<std::int64_t, int, int>> es;
for (int a = 0; a < n; a++) {
for (int b = a + 1; b < n; b++) {
std::int64_t dis = std::max(std::abs(pt[a].first - pt[b].first),
std::abs(pt[a].second - pt[b].second));
es.emplace_back(dis, a, b);
}
}
std::sort(es.begin(), es.end());
atcoder::dsu dsu(n);
std::int64_t cst = 0;
std::vector<std::pair<int, int>> res;
for (auto [w, a, b] : es) {
if (!dsu.same(a, b)) {
dsu.merge(a, b);
cst += w;
res.emplace_back(a, b);
}
}
return {cst, res};
}
} // namespace noya