matrix_tree.hpp¶
用矩阵树定理计算无向生成树或有向根生成树数量；适合对所有连通树形选择计数。
Complexity: Time: \(O(V^{3} + E \alpha (V))\) undirected; \(O(V^{3} + E)\) directed. Space: \(O(V^{2})\). Field operations are assumed constant time.
AC 记录：counting_spanning_tree_directed, counting_spanning_tree_undirected。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(V^3 + E alpha(V)) undirected; O(V^3 + E) directed.
/// Space: O(V^2). Field operations are assumed constant time.
#include "noya/math/linear_algebra.hpp"
#include "atcoder/dsu.hpp"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
/// @brief Count weighted spanning trees of an undirected graph over a field
/// using the Matrix-Tree theorem. Loops are ignored; parallel edges are labeled.
template <class T>
T spanning_tree_count(int n, const std::vector<std::tuple<int, int, T>> &es) {
assert(n >= 0);
if (n <= 1) {
return T(1);
}
if (es.size() < std::size_t(n - 1)) return T(0);
// All principal cofactors agree. Deleting a high-degree vertex can reduce
// elimination fill; zero-weight edges and loops do not affect connectivity.
std::vector<std::size_t> degree(n);
atcoder::dsu connected(n);
int components = n;
for (const auto &[a, b, w] : es) {
assert(0 <= a && a < n && 0 <= b && b < n);
if (a != b && w != T(0)) {
++degree[a];
++degree[b];
if (components > 1 && !connected.same(a, b)) {
connected.merge(a, b);
--components;
}
}
}
if (components != 1) return T(0);
int root = n - 1;
for (int u = 0; u < n; ++u)
if (degree[u] > degree[root]) root = u;
std::vector<std::vector<T>> mat(n - 1, std::vector<T>(n - 1));
for (const auto &[u, v, w] : es) {
if (u == v) continue;
int a = u - (u > root), b = v - (v > root);
if (u != root) {
mat[a][a] += w;
}
if (v != root) {
mat[b][b] += w;
}
if (u != root && v != root) {
mat[a][b] -= w;
mat[b][a] -= w;
}
}
return determinant(std::move(mat));
}
/// @brief Count weighted directed spanning trees whose edges point toward root
/// over a field, using the directed Matrix-Tree theorem.
template <class T>
T in_arborescence_count(int n, int rt,
const std::vector<std::tuple<int, int, T>> &es) {
assert(n > 0);
assert(0 <= rt && rt < n);
if (n == 1) {
return T(1);
}
if (es.size() < std::size_t(n - 1)) return T(0);
std::vector<bool> outgoing(n);
outgoing[rt] = true;
int missing = n - 1;
for (const auto &[u, to, w] : es) {
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to && w != T(0) && !outgoing[u]) {
outgoing[u] = true;
if (--missing == 0) break;
}
}
if (missing) return T(0);
// Construct only the minor: outgoing edges of the deleted root do not count.
std::vector<std::vector<T>> mat(n - 1, std::vector<T>(n - 1));
for (const auto &[u, to, w] : es) {
assert(0 <= u && u < n && 0 <= to && to < n);
if (u == rt || u == to) continue;
int a = u - (u > rt), b = to - (to > rt);
mat[a][a] += w;
if (to != rt) mat[a][b] -= w;
}
return determinant(std::move(mat));
}
} // namespace noya
// NOYA_MATRIX_TREE_HPP
#ifndef NOYA_MATRIX_TREE_HPP
#define NOYA_MATRIX_TREE_HPP 1
/// @complexity Time: O(V^3 + E alpha(V)) undirected; O(V^3 + E) directed.
/// Space: O(V^2). Field operations are assumed constant time.
#include "noya/math/linear_algebra.hpp"
#include "atcoder/dsu.hpp"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
/// @brief Count weighted spanning trees of an undirected graph over a field
/// using the Matrix-Tree theorem. Loops are ignored; parallel edges are labeled.
template <class T>
T spanning_tree_count(int n, const std::vector<std::tuple<int, int, T>> &es) {
assert(n >= 0);
if (n <= 1) {
return T(1);
}
if (es.size() < std::size_t(n - 1)) return T(0);
// All principal cofactors agree. Deleting a high-degree vertex can reduce
// elimination fill; zero-weight edges and loops do not affect connectivity.
std::vector<std::size_t> degree(n);
atcoder::dsu connected(n);
int components = n;
for (const auto &[a, b, w] : es) {
assert(0 <= a && a < n && 0 <= b && b < n);
if (a != b && w != T(0)) {
++degree[a];
++degree[b];
if (components > 1 && !connected.same(a, b)) {
connected.merge(a, b);
--components;
}
}
}
if (components != 1) return T(0);
int root = n - 1;
for (int u = 0; u < n; ++u)
if (degree[u] > degree[root]) root = u;
std::vector<std::vector<T>> mat(n - 1, std::vector<T>(n - 1));
for (const auto &[u, v, w] : es) {
if (u == v) continue;
int a = u - (u > root), b = v - (v > root);
if (u != root) {
mat[a][a] += w;
}
if (v != root) {
mat[b][b] += w;
}
if (u != root && v != root) {
mat[a][b] -= w;
mat[b][a] -= w;
}
}
return determinant(std::move(mat));
}
/// @brief Count weighted directed spanning trees whose edges point toward root
/// over a field, using the directed Matrix-Tree theorem.
template <class T>
T in_arborescence_count(int n, int rt,
const std::vector<std::tuple<int, int, T>> &es) {
assert(n > 0);
assert(0 <= rt && rt < n);
if (n == 1) {
return T(1);
}
if (es.size() < std::size_t(n - 1)) return T(0);
std::vector<bool> outgoing(n);
outgoing[rt] = true;
int missing = n - 1;
for (const auto &[u, to, w] : es) {
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to && w != T(0) && !outgoing[u]) {
outgoing[u] = true;
if (--missing == 0) break;
}
}
if (missing) return T(0);
// Construct only the minor: outgoing edges of the deleted root do not count.
std::vector<std::vector<T>> mat(n - 1, std::vector<T>(n - 1));
for (const auto &[u, to, w] : es) {
assert(0 <= u && u < n && 0 <= to && to < n);
if (u == rt || u == to) continue;
int a = u - (u > rt), b = to - (to > rt);
mat[a][a] += w;
if (to != rt) mat[a][b] -= w;
}
return determinant(std::move(mat));
}
} // namespace noya
#endif // NOYA_MATRIX_TREE_HPP
/// @complexity Time: O(V^3 + E alpha(V)) undirected; O(V^3 + E) directed.
/// Space: O(V^2). Field operations are assumed constant time.
/// @complexity Time: O(nm min(n,m)) elimination and O(n^3) square determinant.
/// Solving also takes O(m(m-r)) to construct the nullspace basis.
/// Space: O(nm) for rank/determinant; O(nm + m(m-r)) for solving.
/// n and m are the row and column counts, and r is the rank.
/// Field operations are counted as O(1).
#include <algorithm>
#include <cassert>
#include <type_traits>
#include <utility>
#include <vector>
namespace noya {
/// @brief Exact zero predicate used by elimination routines by default.
template <class T> struct exact_zero {
bool operator()(const T &val) const { return val == T{}; }
};
/// @brief Consistency flag, one solution, nullspace basis, and pivot columns.
template <class T> struct linear_system_solution {
bool ok = false;
std::vector<T> sol;
std::vector<std::vector<T>> ker;
std::vector<int> pc;
};
namespace linear_algebra_internal {
// Direct division avoids overflowing the reciprocal of a tiny floating pivot.
template <class T>
T pivot_ratio(const T &val, const T &piv, T &inv, bool &ready) {
if constexpr (std::is_floating_point_v<T>) return val / piv;
if (!ready) {
inv = T(1) / piv;
ready = true;
}
return val * inv;
}
template <class T> int column_count(const std::vector<std::vector<T>> &mat) {
if (mat.empty()) {
return 0;
}
int cs = int(mat[0].size());
for (const auto &row : mat) {
assert(int(row.size()) == cs);
}
return cs;
}
} // namespace linear_algebra_internal
/// @brief Compute matrix rank over a field.
template <class T, class IsZero = exact_zero<T>>
int matrix_rank(std::vector<std::vector<T>> mat, IsZero iz = {}) {
int rs = int(mat.size());
int cs = linear_algebra_internal::column_count(mat);
int ran = 0;
for (int col = 0; col < cs && ran < rs; col++) {
int piv = ran;
while (piv < rs && iz(mat[piv][col])) {
piv++;
}
if (piv == rs) {
continue;
}
std::swap(mat[piv], mat[ran]);
T inv{};
bool ready = false;
for (int row = ran + 1; row < rs; row++) {
if (iz(mat[row][col])) {
continue;
}
T r = linear_algebra_internal::pivot_ratio(
mat[row][col], mat[ran][col], inv, ready);
for (int j = col; j < cs; j++) {
mat[row][j] -= r * mat[ran][j];
}
}
ran++;
}
return ran;
}
/// @brief Compute the determinant of a square matrix over a field.
template <class T, class IsZero = exact_zero<T>>
T determinant(std::vector<std::vector<T>> mat, IsZero iz = {}) {
int n = int(mat.size());
assert(linear_algebra_internal::column_count(mat) == n);
T res = T(1);
for (int col = 0; col < n; col++) {
int piv = col;
while (piv < n && iz(mat[piv][col])) {
piv++;
}
if (piv == n) {
return T{};
}
if (piv != col) {
std::swap(mat[piv], mat[col]);
res = -res;
}
T pv = mat[col][col];
res *= pv;
T inv{};
bool ready = false;
for (int row = col + 1; row < n; row++) {
if (iz(mat[row][col])) {
continue;
}
T r = linear_algebra_internal::pivot_ratio(mat[row][col], pv, inv, ready);
for (int j = col; j < n; j++) {
mat[row][j] -= r * mat[col][j];
}
}
}
return res;
}
/// @brief Solve A*x=b and return one solution plus a basis of the nullspace.
template <class T, class IsZero = exact_zero<T>>
linear_system_solution<T> solve_linear(std::vector<std::vector<T>> mat,
std::vector<T> rhs, IsZero iz = {}) {
int rs = int(mat.size());
assert(int(rhs.size()) == rs);
int cs = linear_algebra_internal::column_count(mat);
std::vector<int> pc;
int ran = 0;
for (int col = 0; col < cs && ran < rs; col++) {
int piv = ran;
while (piv < rs && iz(mat[piv][col])) {
piv++;
}
if (piv == rs) {
continue;
}
std::swap(mat[piv], mat[ran]);
std::swap(rhs[piv], rhs[ran]);
T inv = T(1) / mat[ran][col];
for (int j = col; j < cs; j++) {
mat[ran][j] *= inv;
}
rhs[ran] *= inv;
for (int row = 0; row < rs; row++) {
if (row == ran || iz(mat[row][col])) {
continue;
}
T r = mat[row][col];
for (int j = col; j < cs; j++) {
mat[row][j] -= r * mat[ran][j];
}
rhs[row] -= r * rhs[ran];
}
pc.push_back(col);
ran++;
}
for (int row = ran; row < rs; row++) {
bool az = true;
for (int col = 0; col < cs; col++) {
az &= iz(mat[row][col]);
}
if (az && !iz(rhs[row])) {
return {};
}
}
linear_system_solution<T> res;
res.ok = true;
res.sol.assign(cs, T{});
res.pc = pc;
std::vector<bool> ip(cs);
for (int row = 0; row < ran; row++) {
int col = pc[row];
ip[col] = true;
res.sol[col] = rhs[row];
}
for (int fc = 0; fc < cs; fc++) {
if (ip[fc]) {
continue;
}
std::vector<T> bv(cs, T{});
bv[fc] = T(1);
for (int row = 0; row < ran; row++) {
bv[pc[row]] = -mat[row][fc];
}
res.ker.push_back(std::move(bv));
}
return res;
}
} // namespace noya
#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 <cstddef>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
/// @brief Count weighted spanning trees of an undirected graph over a field
/// using the Matrix-Tree theorem. Loops are ignored; parallel edges are labeled.
template <class T>
T spanning_tree_count(int n, const std::vector<std::tuple<int, int, T>> &es) {
assert(n >= 0);
if (n <= 1) {
return T(1);
}
if (es.size() < std::size_t(n - 1)) return T(0);
// All principal cofactors agree. Deleting a high-degree vertex can reduce
// elimination fill; zero-weight edges and loops do not affect connectivity.
std::vector<std::size_t> degree(n);
atcoder::dsu connected(n);
int components = n;
for (const auto &[a, b, w] : es) {
assert(0 <= a && a < n && 0 <= b && b < n);
if (a != b && w != T(0)) {
++degree[a];
++degree[b];
if (components > 1 && !connected.same(a, b)) {
connected.merge(a, b);
--components;
}
}
}
if (components != 1) return T(0);
int root = n - 1;
for (int u = 0; u < n; ++u)
if (degree[u] > degree[root]) root = u;
std::vector<std::vector<T>> mat(n - 1, std::vector<T>(n - 1));
for (const auto &[u, v, w] : es) {
if (u == v) continue;
int a = u - (u > root), b = v - (v > root);
if (u != root) {
mat[a][a] += w;
}
if (v != root) {
mat[b][b] += w;
}
if (u != root && v != root) {
mat[a][b] -= w;
mat[b][a] -= w;
}
}
return determinant(std::move(mat));
}
/// @brief Count weighted directed spanning trees whose edges point toward root
/// over a field, using the directed Matrix-Tree theorem.
template <class T>
T in_arborescence_count(int n, int rt,
const std::vector<std::tuple<int, int, T>> &es) {
assert(n > 0);
assert(0 <= rt && rt < n);
if (n == 1) {
return T(1);
}
if (es.size() < std::size_t(n - 1)) return T(0);
std::vector<bool> outgoing(n);
outgoing[rt] = true;
int missing = n - 1;
for (const auto &[u, to, w] : es) {
assert(0 <= u && u < n);
assert(0 <= to && to < n);
if (u != to && w != T(0) && !outgoing[u]) {
outgoing[u] = true;
if (--missing == 0) break;
}
}
if (missing) return T(0);
// Construct only the minor: outgoing edges of the deleted root do not count.
std::vector<std::vector<T>> mat(n - 1, std::vector<T>(n - 1));
for (const auto &[u, to, w] : es) {
assert(0 <= u && u < n && 0 <= to && to < n);
if (u == rt || u == to) continue;
int a = u - (u > rt), b = to - (to > rt);
mat[a][a] += w;
if (to != rt) mat[a][b] -= w;
}
return determinant(std::move(mat));
}
} // namespace noya
// NOYA_MATRIX_TREE_HPP