matrix_tree.hpp¶
Count weighted spanning trees of an undirected graph over a field using the Matrix-Tree theorem in O(n^3) time.
Verified by counting_spanning_tree_directed, counting_spanning_tree_undirected.
用矩阵树定理计算无向生成树或有向根生成树数量；适合对所有连通树形选择计数。
Implementation¶
#ifndef NOYA_MATRIX_TREE_HPP
#define NOYA_MATRIX_TREE_HPP 1
/// @complexity Time: O(V^3 + E) field operations.
/// Space: O(V^2).
#include "noya/linear_algebra.hpp"
#include <cassert>
#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 in O(n^3) time.
template <class T>
T spanning_tree_count(int n,
const std::vector<std::tuple<int, int, T>> &edges) {
assert(n >= 0);
if (n <= 1) {
return T(1);
}
std::vector<std::vector<T>> minor(n - 1, std::vector<T>(n - 1));
for (const auto &[first, second, weight] : edges) {
assert(0 <= first && first < n);
assert(0 <= second && second < n);
if (first == second) {
continue;
}
if (first < n - 1) {
minor[first][first] += weight;
}
if (second < n - 1) {
minor[second][second] += weight;
}
if (first < n - 1 && second < n - 1) {
minor[first][second] -= weight;
minor[second][first] -= weight;
}
}
return determinant(std::move(minor));
}
/// @brief Count weighted directed spanning trees whose edges point toward root
/// over a field, using the directed Matrix-Tree theorem in O(n^3) time.
template <class T>
T in_arborescence_count(int n, int root,
const std::vector<std::tuple<int, int, T>> &edges) {
assert(n > 0);
assert(0 <= root && root < n);
if (n == 1) {
return T(1);
}
std::vector<std::vector<T>> laplacian(n, std::vector<T>(n));
for (const auto &[from, to, weight] : edges) {
assert(0 <= from && from < n);
assert(0 <= to && to < n);
if (from == to) {
continue;
}
laplacian[from][from] += weight;
laplacian[from][to] -= weight;
}
std::vector<std::vector<T>> minor;
minor.reserve(n - 1);
for (int row = 0; row < n; row++) {
if (row == root) {
continue;
}
minor.emplace_back();
minor.back().reserve(n - 1);
for (int column = 0; column < n; column++) {
if (column != root) {
minor.back().push_back(laplacian[row][column]);
}
}
}
return determinant(std::move(minor));
}
} // namespace noya
#endif // NOYA_MATRIX_TREE_HPP
#include <algorithm>
#include <cassert>
#include <tuple>
#include <utility>
#include <vector>
/// @complexity Time: O(V^3 + E) field operations.
/// Space: O(V^2).
/// @complexity Time: O(rows * columns * min(rows,columns)) elimination; O(n^3) square determinant/inverse.
/// Space: O(rows * columns).
namespace noya {
/// @brief Exact zero predicate used by elimination routines by default.
template <class T> struct exact_zero {
bool operator()(const T &value) const { return value == T{}; }
};
/// @brief Consistency flag, one solution, nullspace basis, and pivot columns.
template <class T> struct linear_system_solution {
bool consistent = false;
std::vector<T> solution;
std::vector<std::vector<T>> nullspace_basis;
std::vector<int> pivot_columns;
};
namespace linear_algebra_internal {
template <class T> int column_count(const std::vector<std::vector<T>> &matrix) {
if (matrix.empty()) {
return 0;
}
int columns = int(matrix[0].size());
for (const auto &row : matrix) {
assert(int(row.size()) == columns);
}
return columns;
}
} // 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>> matrix, IsZero is_zero = {}) {
int rows = int(matrix.size());
int columns = linear_algebra_internal::column_count(matrix);
int rank = 0;
for (int column = 0; column < columns && rank < rows; column++) {
int pivot = rank;
while (pivot < rows && is_zero(matrix[pivot][column])) {
pivot++;
}
if (pivot == rows) {
continue;
}
std::swap(matrix[pivot], matrix[rank]);
for (int row = rank + 1; row < rows; row++) {
if (is_zero(matrix[row][column])) {
continue;
}
T ratio = matrix[row][column] / matrix[rank][column];
for (int j = column; j < columns; j++) {
matrix[row][j] -= ratio * matrix[rank][j];
}
}
rank++;
}
return rank;
}
/// @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>> matrix, IsZero is_zero = {}) {
int n = int(matrix.size());
assert(linear_algebra_internal::column_count(matrix) == n);
T result = T(1);
for (int column = 0; column < n; column++) {
int pivot = column;
while (pivot < n && is_zero(matrix[pivot][column])) {
pivot++;
}
if (pivot == n) {
return T{};
}
if (pivot != column) {
std::swap(matrix[pivot], matrix[column]);
result = -result;
}
T pivot_value = matrix[column][column];
result *= pivot_value;
for (int row = column + 1; row < n; row++) {
if (is_zero(matrix[row][column])) {
continue;
}
T ratio = matrix[row][column] / pivot_value;
for (int j = column; j < n; j++) {
matrix[row][j] -= ratio * matrix[column][j];
}
}
}
return result;
}
/// @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>> matrix,
std::vector<T> right_hand_side,
IsZero is_zero = {}) {
int rows = int(matrix.size());
assert(int(right_hand_side.size()) == rows);
int columns = linear_algebra_internal::column_count(matrix);
std::vector<int> pivot_columns;
int rank = 0;
for (int column = 0; column < columns && rank < rows; column++) {
int pivot = rank;
while (pivot < rows && is_zero(matrix[pivot][column])) {
pivot++;
}
if (pivot == rows) {
continue;
}
std::swap(matrix[pivot], matrix[rank]);
std::swap(right_hand_side[pivot], right_hand_side[rank]);
T inverse = T(1) / matrix[rank][column];
for (int j = column; j < columns; j++) {
matrix[rank][j] *= inverse;
}
right_hand_side[rank] *= inverse;
for (int row = 0; row < rows; row++) {
if (row == rank || is_zero(matrix[row][column])) {
continue;
}
T ratio = matrix[row][column];
for (int j = column; j < columns; j++) {
matrix[row][j] -= ratio * matrix[rank][j];
}
right_hand_side[row] -= ratio * right_hand_side[rank];
}
pivot_columns.push_back(column);
rank++;
}
for (int row = rank; row < rows; row++) {
bool all_zero = true;
for (int column = 0; column < columns; column++) {
all_zero &= is_zero(matrix[row][column]);
}
if (all_zero && !is_zero(right_hand_side[row])) {
return {};
}
}
linear_system_solution<T> result;
result.consistent = true;
result.solution.assign(columns, T{});
result.pivot_columns = pivot_columns;
std::vector<bool> is_pivot(columns);
for (int row = 0; row < rank; row++) {
int column = pivot_columns[row];
is_pivot[column] = true;
result.solution[column] = right_hand_side[row];
}
for (int free_column = 0; free_column < columns; free_column++) {
if (is_pivot[free_column]) {
continue;
}
std::vector<T> basis_vector(columns, T{});
basis_vector[free_column] = T(1);
for (int row = 0; row < rank; row++) {
basis_vector[pivot_columns[row]] = -matrix[row][free_column];
}
result.nullspace_basis.push_back(std::move(basis_vector));
}
return result;
}
} // namespace noya
namespace noya {
/// @brief Count weighted spanning trees of an undirected graph over a field
/// using the Matrix-Tree theorem in O(n^3) time.
template <class T>
T spanning_tree_count(int n,
const std::vector<std::tuple<int, int, T>> &edges) {
assert(n >= 0);
if (n <= 1) {
return T(1);
}
std::vector<std::vector<T>> minor(n - 1, std::vector<T>(n - 1));
for (const auto &[first, second, weight] : edges) {
assert(0 <= first && first < n);
assert(0 <= second && second < n);
if (first == second) {
continue;
}
if (first < n - 1) {
minor[first][first] += weight;
}
if (second < n - 1) {
minor[second][second] += weight;
}
if (first < n - 1 && second < n - 1) {
minor[first][second] -= weight;
minor[second][first] -= weight;
}
}
return determinant(std::move(minor));
}
/// @brief Count weighted directed spanning trees whose edges point toward root
/// over a field, using the directed Matrix-Tree theorem in O(n^3) time.
template <class T>
T in_arborescence_count(int n, int root,
const std::vector<std::tuple<int, int, T>> &edges) {
assert(n > 0);
assert(0 <= root && root < n);
if (n == 1) {
return T(1);
}
std::vector<std::vector<T>> laplacian(n, std::vector<T>(n));
for (const auto &[from, to, weight] : edges) {
assert(0 <= from && from < n);
assert(0 <= to && to < n);
if (from == to) {
continue;
}
laplacian[from][from] += weight;
laplacian[from][to] -= weight;
}
std::vector<std::vector<T>> minor;
minor.reserve(n - 1);
for (int row = 0; row < n; row++) {
if (row == root) {
continue;
}
minor.emplace_back();
minor.back().reserve(n - 1);
for (int column = 0; column < n; column++) {
if (column != root) {
minor.back().push_back(laplacian[row][column]);
}
}
}
return determinant(std::move(minor));
}
} // namespace noya