matrix/linear-equation.hpp

• View this file on GitHub
• Last update: 2024-05-03 21:06:15+09:00
• Include: #include "matrix/linear-equation.hpp"

Depends on

• matrix/gauss-elimination.hpp

Required by

• P-recursiveの高速計算 (fps/find-p-recursive.hpp)

Verified with

• verify/verify-aoj-other/aoj-2171-bigrational.test.cpp
• verify/verify-aoj-other/aoj-2171.test.cpp
• verify/verify-unit-test/gauss-elimination.test.cpp
• verify/verify-unit-test/p-recursive.test.cpp
• verify/verify-yosupo-fps/yosupo-factorial-p-recursive.test.cpp
• verify/verify-yosupo-math/yosupo-linear-equation-2.test.cpp
• verify/verify-yuki/yuki-1533.test.cpp

Code

#pragma once

#include "gauss-elimination.hpp"

// 解が存在する場合は, 解が v + C_1 w_1 + ... + C_k w_k と表せるとして
// (v, w_1, ..., w_k) を返す
// 解が存在しない場合は空のベクトルを返す
//
// double や Rational でも動くはず？(未検証)
template <typename T>
vector<vector<T>> LinearEquation(vector<vector<T>> a, vector<T> b) {
  int H = a.size(), W = a[0].size();
  for (int i = 0; i < H; i++) a[i].push_back(b[i]);
  auto p = GaussElimination(a, W, true);
  int rank = p.first;
  for (int i = rank; i < H; ++i) {
    if (a[i][W] != 0) return vector<vector<T>>{};
  }
  vector<vector<T>> res(1, vector<T>(W));
  vector<int> pivot(W, -1);
  for (int i = 0, j = 0; i < rank; ++i) {
    while (a[i][j] == 0) ++j;
    res[0][j] = a[i][W], pivot[j] = i;
  }
  for (int j = 0; j < W; ++j) {
    if (pivot[j] == -1) {
      vector<T> x(W);
      x[j] = 1;
      for (int k = 0; k < j; ++k) {
        if (pivot[k] != -1) x[k] = -a[pivot[k]][j];
      }
      res.push_back(x);
    }
  }
  return res;
}

#line 2 "matrix/linear-equation.hpp"

#line 2 "matrix/gauss-elimination.hpp"

#include <utility>
#include <vector>
using namespace std;

// {rank, det(非正方行列の場合は未定義)} を返す
// 型が double や Rational でも動くはず？(未検証)
//
// pivot 候補 : [0, pivot_end)
template <typename T>
std::pair<int, T> GaussElimination(vector<vector<T>> &a, int pivot_end = -1,
                                   bool diagonalize = false) {
  if (a.empty()) return {0, 1};
  int H = a.size(), W = a[0].size(), rank = 0;
  if (pivot_end == -1) pivot_end = W;
  T det = 1;
  for (int j = 0; j < pivot_end; j++) {
    int idx = -1;
    for (int i = rank; i < H; i++) {
      if (a[i][j] != T(0)) {
        idx = i;
        break;
      }
    }
    if (idx == -1) {
      det = 0;
      continue;
    }
    if (rank != idx) det = -det, swap(a[rank], a[idx]);
    det *= a[rank][j];
    if (diagonalize && a[rank][j] != T(1)) {
      T coeff = T(1) / a[rank][j];
      for (int k = j; k < W; k++) a[rank][k] *= coeff;
    }
    int is = diagonalize ? 0 : rank + 1;
    for (int i = is; i < H; i++) {
      if (i == rank) continue;
      if (a[i][j] != T(0)) {
        T coeff = a[i][j] / a[rank][j];
        for (int k = j; k < W; k++) a[i][k] -= a[rank][k] * coeff;
      }
    }
    rank++;
  }
  return make_pair(rank, det);
}
#line 4 "matrix/linear-equation.hpp"

// 解が存在する場合は, 解が v + C_1 w_1 + ... + C_k w_k と表せるとして
// (v, w_1, ..., w_k) を返す
// 解が存在しない場合は空のベクトルを返す
//
// double や Rational でも動くはず？(未検証)
template <typename T>
vector<vector<T>> LinearEquation(vector<vector<T>> a, vector<T> b) {
  int H = a.size(), W = a[0].size();
  for (int i = 0; i < H; i++) a[i].push_back(b[i]);
  auto p = GaussElimination(a, W, true);
  int rank = p.first;
  for (int i = rank; i < H; ++i) {
    if (a[i][W] != 0) return vector<vector<T>>{};
  }
  vector<vector<T>> res(1, vector<T>(W));
  vector<int> pivot(W, -1);
  for (int i = 0, j = 0; i < rank; ++i) {
    while (a[i][j] == 0) ++j;
    res[0][j] = a[i][W], pivot[j] = i;
  }
  for (int j = 0; j < W; ++j) {
    if (pivot[j] == -1) {
      vector<T> x(W);
      x[j] = 1;
      for (int k = 0; k < j; ++k) {
        if (pivot[k] != -1) x[k] = -a[pivot[k]][j];
      }
      res.push_back(x);
    }
  }
  return res;
}

Back to top page