Skip to content

matrix.hpp

SECTIONMath INCLUDEnoya/matrix.hpp

Dense matrix multiplication in O(nmk).

Verified by inverse_matrix, matrix_product, pow_of_matrix.

\[ \displaystyle C=AB \]

Implementation

View on GitHub

#ifndef NOYA_MATRIX_HPP
#define NOYA_MATRIX_HPP 1

/// @complexity Time: O(nmk) multiplication and O(n^3) square inversion.
/// Space: O(nm) result plus O(n^2) inversion workspace.

#include "noya/linear_algebra.hpp"

#include <cassert>
#include <cstdint>
#include <optional>
#include <utility>
#include <vector>

namespace noya {

template <class T> using matrix = std::vector<std::vector<T>>;

template <class T> matrix<T> identity_matrix(int n) {
  assert(n >= 0);
  matrix<T> result(n, std::vector<T>(n));
  for (int index = 0; index < n; index++) {
    result[index][index] = T(1);
  }
  return result;
}

/// @brief Dense matrix multiplication in O(nmk).
template <class T>
matrix<T> matrix_multiply(const matrix<T> &first, const matrix<T> &second) {
  int rows = int(first.size());
  int middle = rows == 0 ? int(second.size()) : int(first[0].size());
  for (const auto &row : first) {
    assert(int(row.size()) == middle);
  }
  assert(int(second.size()) == middle);
  int columns = middle == 0 ? 0 : int(second[0].size());
  for (const auto &row : second) {
    assert(int(row.size()) == columns);
  }
  matrix<T> result(rows, std::vector<T>(columns));
  for (int row = 0; row < rows; row++) {
    for (int index = 0; index < middle; index++) {
      for (int column = 0; column < columns; column++) {
        result[row][column] += first[row][index] * second[index][column];
      }
    }
  }
  return result;
}

template <class T>
matrix<T> matrix_power(matrix<T> value, std::uint64_t exponent) {
  int n = int(value.size());
  for (const auto &row : value) {
    assert(int(row.size()) == n);
  }
  matrix<T> result = identity_matrix<T>(n);
  while (exponent > 0) {
    if (exponent & 1) {
      result = matrix_multiply(result, value);
    }
    value = matrix_multiply(value, value);
    exponent >>= 1;
  }
  return result;
}

/// @brief Multiply matrices of static modular integers, transposing the right
/// operand and reducing one wide accumulator per output entry.
template <class Mint>
matrix<Mint> matrix_multiply_mod(const matrix<Mint> &first,
                                 const matrix<Mint> &second) {
  int rows = int(first.size());
  int middle = rows == 0 ? int(second.size()) : int(first[0].size());
  int columns = middle == 0 ? 0 : int(second[0].size());
  matrix<unsigned int> transposed(columns,
                                  std::vector<unsigned int>(middle));
  for (int index = 0; index < middle; index++) {
    for (int column = 0; column < columns; column++) {
      transposed[column][index] = second[index][column].val();
    }
  }
  matrix<Mint> result(rows, std::vector<Mint>(columns));
  for (int row = 0; row < rows; row++) {
    for (int column = 0; column < columns; column++) {
      unsigned __int128 sum = 0;
      for (int index = 0; index < middle; index++) {
        sum += static_cast<std::uint64_t>(first[row][index].val()) *
               transposed[column][index];
      }
      result[row][column] = Mint::raw(unsigned(sum % Mint::mod()));
    }
  }
  return result;
}

/// @brief Binary exponentiation specialized for dense modular matrices.
template <class Mint>
matrix<Mint> matrix_power_mod(matrix<Mint> value, std::uint64_t exponent) {
  int n = int(value.size());
  matrix<Mint> result = identity_matrix<Mint>(n);
  while (exponent > 0) {
    if (exponent & 1) {
      result = matrix_multiply_mod(result, value);
    }
    exponent >>= 1;
    if (exponent != 0) {
      value = matrix_multiply_mod(value, value);
    }
  }
  return result;
}

/// @brief Invert a square matrix over a field, returning nullopt when singular.
template <class T, class IsZero = exact_zero<T>>
std::optional<matrix<T>> matrix_inverse(matrix<T> value, IsZero is_zero = {}) {
  int n = int(value.size());
  for (const auto &row : value) {
    assert(int(row.size()) == n);
  }
  matrix<T> inverse = identity_matrix<T>(n);
  for (int column = 0; column < n; column++) {
    int pivot = column;
    while (pivot < n && is_zero(value[pivot][column])) {
      pivot++;
    }
    if (pivot == n) {
      return std::nullopt;
    }
    std::swap(value[pivot], value[column]);
    std::swap(inverse[pivot], inverse[column]);
    T scale = T(1) / value[column][column];
    for (int index = 0; index < n; index++) {
      value[column][index] *= scale;
      inverse[column][index] *= scale;
    }
    for (int row = 0; row < n; row++) {
      if (row == column || is_zero(value[row][column])) {
        continue;
      }
      T ratio = value[row][column];
      for (int index = 0; index < n; index++) {
        value[row][index] -= ratio * value[column][index];
        inverse[row][index] -= ratio * inverse[column][index];
      }
    }
  }
  return inverse;
}

} // namespace noya

#endif // NOYA_MATRIX_HPP
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <optional>
#include <utility>
#include <vector>

/// @complexity Time: O(nmk) multiplication and O(n^3) square inversion.
/// Space: O(nm) result plus O(n^2) inversion workspace.

/// @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 {

template <class T> using matrix = std::vector<std::vector<T>>;

template <class T> matrix<T> identity_matrix(int n) {
  assert(n >= 0);
  matrix<T> result(n, std::vector<T>(n));
  for (int index = 0; index < n; index++) {
    result[index][index] = T(1);
  }
  return result;
}

/// @brief Dense matrix multiplication in O(nmk).
template <class T>
matrix<T> matrix_multiply(const matrix<T> &first, const matrix<T> &second) {
  int rows = int(first.size());
  int middle = rows == 0 ? int(second.size()) : int(first[0].size());
  for (const auto &row : first) {
    assert(int(row.size()) == middle);
  }
  assert(int(second.size()) == middle);
  int columns = middle == 0 ? 0 : int(second[0].size());
  for (const auto &row : second) {
    assert(int(row.size()) == columns);
  }
  matrix<T> result(rows, std::vector<T>(columns));
  for (int row = 0; row < rows; row++) {
    for (int index = 0; index < middle; index++) {
      for (int column = 0; column < columns; column++) {
        result[row][column] += first[row][index] * second[index][column];
      }
    }
  }
  return result;
}

template <class T>
matrix<T> matrix_power(matrix<T> value, std::uint64_t exponent) {
  int n = int(value.size());
  for (const auto &row : value) {
    assert(int(row.size()) == n);
  }
  matrix<T> result = identity_matrix<T>(n);
  while (exponent > 0) {
    if (exponent & 1) {
      result = matrix_multiply(result, value);
    }
    value = matrix_multiply(value, value);
    exponent >>= 1;
  }
  return result;
}

/// @brief Multiply matrices of static modular integers, transposing the right
/// operand and reducing one wide accumulator per output entry.
template <class Mint>
matrix<Mint> matrix_multiply_mod(const matrix<Mint> &first,
                                 const matrix<Mint> &second) {
  int rows = int(first.size());
  int middle = rows == 0 ? int(second.size()) : int(first[0].size());
  int columns = middle == 0 ? 0 : int(second[0].size());
  matrix<unsigned int> transposed(columns,
                                  std::vector<unsigned int>(middle));
  for (int index = 0; index < middle; index++) {
    for (int column = 0; column < columns; column++) {
      transposed[column][index] = second[index][column].val();
    }
  }
  matrix<Mint> result(rows, std::vector<Mint>(columns));
  for (int row = 0; row < rows; row++) {
    for (int column = 0; column < columns; column++) {
      unsigned __int128 sum = 0;
      for (int index = 0; index < middle; index++) {
        sum += static_cast<std::uint64_t>(first[row][index].val()) *
               transposed[column][index];
      }
      result[row][column] = Mint::raw(unsigned(sum % Mint::mod()));
    }
  }
  return result;
}

/// @brief Binary exponentiation specialized for dense modular matrices.
template <class Mint>
matrix<Mint> matrix_power_mod(matrix<Mint> value, std::uint64_t exponent) {
  int n = int(value.size());
  matrix<Mint> result = identity_matrix<Mint>(n);
  while (exponent > 0) {
    if (exponent & 1) {
      result = matrix_multiply_mod(result, value);
    }
    exponent >>= 1;
    if (exponent != 0) {
      value = matrix_multiply_mod(value, value);
    }
  }
  return result;
}

/// @brief Invert a square matrix over a field, returning nullopt when singular.
template <class T, class IsZero = exact_zero<T>>
std::optional<matrix<T>> matrix_inverse(matrix<T> value, IsZero is_zero = {}) {
  int n = int(value.size());
  for (const auto &row : value) {
    assert(int(row.size()) == n);
  }
  matrix<T> inverse = identity_matrix<T>(n);
  for (int column = 0; column < n; column++) {
    int pivot = column;
    while (pivot < n && is_zero(value[pivot][column])) {
      pivot++;
    }
    if (pivot == n) {
      return std::nullopt;
    }
    std::swap(value[pivot], value[column]);
    std::swap(inverse[pivot], inverse[column]);
    T scale = T(1) / value[column][column];
    for (int index = 0; index < n; index++) {
      value[column][index] *= scale;
      inverse[column][index] *= scale;
    }
    for (int row = 0; row < n; row++) {
      if (row == column || is_zero(value[row][column])) {
        continue;
      }
      T ratio = value[row][column];
      for (int index = 0; index < n; index++) {
        value[row][index] -= ratio * value[column][index];
        inverse[row][index] -= ratio * inverse[column][index];
      }
    }
  }
  return inverse;
}

} // namespace noya