special_spanning_trees.hpp¶
Find a spanning tree minimizing (sum first_weight)*(sum second_weight) by recursive supported-point search on the bicriteria spanning-tree hull.
计算带特殊边类别或选择限制的生成树数量；适合需要在矩阵树定理上加入组合约束的题目。
Implementation¶
#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 first_sum = 0;
std::int64_t second_sum = 0;
std::vector<int> edge_ids;
};
/// @brief Find a spanning tree minimizing (sum first_weight)*(sum
/// second_weight) 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>>
&edges) {
assert(n >= 1);
auto tree = [&](std::int64_t first_coefficient,
std::int64_t second_coefficient) {
std::vector<int> order(edges.size());
std::iota(order.begin(), order.end(), 0);
std::sort(order.begin(), order.end(), [&](int first, int second) {
auto [a1, b1, u1, v1] = edges[first];
auto [a2, b2, u2, v2] = edges[second];
(void)u1;
(void)v1;
(void)u2;
(void)v2;
__int128 value1 =
__int128(first_coefficient) * a1 + __int128(second_coefficient) * b1;
__int128 value2 =
__int128(first_coefficient) * a2 + __int128(second_coefficient) * b2;
return value1 != value2 ? value1 < value2 : first < second;
});
atcoder::dsu dsu(n);
bicriteria_spanning_tree_result result;
for (int id : order) {
auto [first_weight, second_weight, first, second] = edges[id];
assert(0 <= first && first < n);
assert(0 <= second && second < n);
if (!dsu.same(first, second)) {
dsu.merge(first, second);
result.first_sum += first_weight;
result.second_sum += second_weight;
result.edge_ids.push_back(id);
}
}
assert(int(result.edge_ids.size()) == n - 1);
return result;
};
auto first_extreme = tree(1, 0);
auto second_extreme = tree(0, 1);
auto product = [](const bicriteria_spanning_tree_result &value) {
return __int128(value.first_sum) * value.second_sum;
};
auto best = product(first_extreme) < product(second_extreme) ? first_extreme
: second_extreme;
auto search = [&](auto &self, const bicriteria_spanning_tree_result &left,
const bicriteria_spanning_tree_result &right) -> void {
std::int64_t first_coefficient = left.second_sum - right.second_sum;
std::int64_t second_coefficient = right.first_sum - left.first_sum;
auto middle = tree(first_coefficient, second_coefficient);
__int128 line_value = __int128(first_coefficient) * left.first_sum +
__int128(second_coefficient) * left.second_sum;
__int128 middle_value = __int128(first_coefficient) * middle.first_sum +
__int128(second_coefficient) * middle.second_sum;
if (middle_value >= line_value) {
return;
}
if (product(middle) < product(best)) {
best = middle;
}
self(self, left, middle);
self(self, middle, right);
};
search(search, first_extreme, second_extreme);
return best;
}
/// @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>> &points) {
int n = int(points.size());
std::vector<std::tuple<std::int64_t, int, int>> edges;
for (int first = 0; first < n; first++) {
for (int second = first + 1; second < n; second++) {
std::int64_t distance =
std::max(std::abs(points[first].first - points[second].first),
std::abs(points[first].second - points[second].second));
edges.emplace_back(distance, first, second);
}
}
std::sort(edges.begin(), edges.end());
atcoder::dsu dsu(n);
std::int64_t cost = 0;
std::vector<std::pair<int, int>> result;
for (auto [weight, first, second] : edges) {
if (!dsu.same(first, second)) {
dsu.merge(first, second);
cost += weight;
result.emplace_back(first, second);
}
}
return {cost, result};
}
} // 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 first_sum = 0;
std::int64_t second_sum = 0;
std::vector<int> edge_ids;
};
/// @brief Find a spanning tree minimizing (sum first_weight)*(sum
/// second_weight) 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>>
&edges) {
assert(n >= 1);
auto tree = [&](std::int64_t first_coefficient,
std::int64_t second_coefficient) {
std::vector<int> order(edges.size());
std::iota(order.begin(), order.end(), 0);
std::sort(order.begin(), order.end(), [&](int first, int second) {
auto [a1, b1, u1, v1] = edges[first];
auto [a2, b2, u2, v2] = edges[second];
(void)u1;
(void)v1;
(void)u2;
(void)v2;
__int128 value1 =
__int128(first_coefficient) * a1 + __int128(second_coefficient) * b1;
__int128 value2 =
__int128(first_coefficient) * a2 + __int128(second_coefficient) * b2;
return value1 != value2 ? value1 < value2 : first < second;
});
atcoder::dsu dsu(n);
bicriteria_spanning_tree_result result;
for (int id : order) {
auto [first_weight, second_weight, first, second] = edges[id];
assert(0 <= first && first < n);
assert(0 <= second && second < n);
if (!dsu.same(first, second)) {
dsu.merge(first, second);
result.first_sum += first_weight;
result.second_sum += second_weight;
result.edge_ids.push_back(id);
}
}
assert(int(result.edge_ids.size()) == n - 1);
return result;
};
auto first_extreme = tree(1, 0);
auto second_extreme = tree(0, 1);
auto product = [](const bicriteria_spanning_tree_result &value) {
return __int128(value.first_sum) * value.second_sum;
};
auto best = product(first_extreme) < product(second_extreme) ? first_extreme
: second_extreme;
auto search = [&](auto &self, const bicriteria_spanning_tree_result &left,
const bicriteria_spanning_tree_result &right) -> void {
std::int64_t first_coefficient = left.second_sum - right.second_sum;
std::int64_t second_coefficient = right.first_sum - left.first_sum;
auto middle = tree(first_coefficient, second_coefficient);
__int128 line_value = __int128(first_coefficient) * left.first_sum +
__int128(second_coefficient) * left.second_sum;
__int128 middle_value = __int128(first_coefficient) * middle.first_sum +
__int128(second_coefficient) * middle.second_sum;
if (middle_value >= line_value) {
return;
}
if (product(middle) < product(best)) {
best = middle;
}
self(self, left, middle);
self(self, middle, right);
};
search(search, first_extreme, second_extreme);
return best;
}
/// @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>> &points) {
int n = int(points.size());
std::vector<std::tuple<std::int64_t, int, int>> edges;
for (int first = 0; first < n; first++) {
for (int second = first + 1; second < n; second++) {
std::int64_t distance =
std::max(std::abs(points[first].first - points[second].first),
std::abs(points[first].second - points[second].second));
edges.emplace_back(distance, first, second);
}
}
std::sort(edges.begin(), edges.end());
atcoder::dsu dsu(n);
std::int64_t cost = 0;
std::vector<std::pair<int, int>> result;
for (auto [weight, first, second] : edges) {
if (!dsu.same(first, second)) {
dsu.merge(first, second);
cost += weight;
result.emplace_back(first, second);
}
}
return {cost, result};
}
} // namespace noya