Square Matrix
(src/matrix/SquareMatrix.hpp)

• View this file on GitHub
• Last update: 2024-05-24 00:56:59+09:00
• Include: #include "src/matrix/SquareMatrix.hpp"

主としてサイズの固定された行列群を扱う際に std::array を用いることで各処理を高速化したライブラリ。

下表では扱う行列の行数（列数）を $N$ とする。

メンバ関数	効果	時間計算量
size()	行数を返す．	$\mathrm{O}(1)$
identity()	単位行列を返す．	$\mathrm{O}(N^2)$
加算	行列 $A$ に $B$ を加算する．	$\mathrm{O}(N^2)$
減算	行列 $A$ に $B$ を減算する．	$\mathrm{O}(N^2)$
乗算	行列 $A$ に $B$ を減算する．	$\mathrm{O}(N^3)$
スカラー倍	行列 $A$ をスカラー倍する．	$\mathrm{O}(N^2)$
pow(n)	行列 $A$ を $n$ 乗した行列を返す．	$\mathrm{O}(N^3 \log n)$
transpose()	行列 $A$ を転置した行列を返す．	$\mathrm{O}(N^2)$
rank()	行列 $A$ の rank を返す．	$\mathrm{O}(N^3)$
det()	行列 $A$ の determinant を返す．	$\mathrm{O}(N^3)$
inv()	行列 $A$ の逆行列を返す．	$\mathrm{O}(N^3)$

Verified with

• test/yosupo/inverse_matrix.squarematrix.test.cpp
• test/yosupo/matrix_det.squarematrix.test.cpp
• test/yosupo/pow_of_matrix.squarematrix.test.cpp

Code

#pragma once
#include <array>
#include <cassert>
#include <utility>

template <typename T, int N> struct SquareMatrix {
    std::array<std::array<T, N>, N> A;

    SquareMatrix() : A() {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                A[i][j] = T(0);
            }
        }
    }

    int size() const { return N; }

    inline const std::array<T, N>& operator[](int i) const { return A[i]; }

    inline std::array<T, N>& operator[](int i) { return A[i]; }

    static SquareMatrix identity() {
        SquareMatrix res;
        for (int i = 0; i < N; i++) res[i][i] = 1;
        return res;
    }

    SquareMatrix& operator+=(const SquareMatrix& B) {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] += B[i][j];
            }
        }
        return *this;
    }

    SquareMatrix& operator-=(const SquareMatrix& B) {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] -= B[i][j];
            }
        }
        return *this;
    }

    SquareMatrix& operator*=(const SquareMatrix& B) {
        std::array<std::array<T, N>, N> C;
        for (int i = 0; i < N; i++) {
            for (int k = 0; k < N; k++) {
                for (int j = 0; j < N; j++) {
                    C[i][j] += (*this)[i][k] * B[k][j];
                }
            }
        }
        std::swap(A, C);
        return *this;
    }

    SquareMatrix& operator*=(const T& v) {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] *= v;
            }
        }
        return *this;
    }

    SquareMatrix& operator/=(const T& v) {
        T inv = T(1) / v;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] *= inv;
            }
        }
        return *this;
    }

    SquareMatrix operator-() const {
        SquareMatrix res;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                res[i][j] = -(*this)[i][j];
            }
        }
        return res;
    }

    SquareMatrix operator+(const SquareMatrix& B) const { return SquareMatrix(*this) += B; }

    SquareMatrix operator-(const SquareMatrix& B) const { return SquareMatrix(*this) -= B; }

    SquareMatrix operator*(const SquareMatrix& B) const { return SquareMatrix(*this) *= B; }

    SquareMatrix operator*(const T& v) const { return SquareMatrix(*this) *= v; }

    SquareMatrix operator/(const T& v) const { return SquareMatrix(*this) /= v; }

    bool operator==(const SquareMatrix& B) const { return A == B.A; }

    bool operator!=(const SquareMatrix& B) const { return A != B.A; }

    SquareMatrix pow(long long n) const {
        assert(0 <= n);
        SquareMatrix x = *this, r = identity();
        while (n) {
            if (n & 1) r *= x;
            x *= x;
            n >>= 1;
        }
        return r;
    }

    SquareMatrix transpose() const {
        SquareMatrix res;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                res[j][i] = (*this)[i][j];
            }
        }
        return res;
    }

    int rank() const { return SquareMatrix(*this).gauss_jordan().first; }

    T det() const { return SquareMatrix(*this).gauss_jordan().second; }

    SquareMatrix inv() const {
        SquareMatrix B(*this), C = identity();
        for (int j = 0; j < N; j++) {
            int pivot = -1;
            for (int i = j; i < N; i++) {
                if (B[i][j] != T(0)) {
                    pivot = i;
                    break;
                }
            }
            assert(pivot != -1);
            if (pivot != j) {
                std::swap(B[pivot], B[j]);
                std::swap(C[pivot], C[j]);
            }
            {
                T coef = T(1) / B[j][j];
                for (int k = 0; k < N; k++) {
                    B[j][k] *= coef;
                    C[j][k] *= coef;
                }
            }
            for (int i = 0; i < N; i++) {
                if (i == j) continue;
                T coef = B[i][j];
                if (coef == T(0)) continue;
                for (int k = 0; k < N; k++) {
                    B[i][k] -= B[j][k] * coef;
                    C[i][k] -= C[j][k] * coef;
                }
            }
        }
        return C;
    }

    friend std::ostream& operator<<(std::ostream& os, const SquareMatrix& p) {
        os << "[(" << N << " * " << N << " Matrix)";
        os << "\n[columun sums: ";
        for (int j = 0; j < N; j++) {
            T sum = 0;
            for (int i = 0; i < N; i++) sum += p[i][j];
            ;
            os << sum << (j + 1 < N ? "," : "");
        }
        os << "]";
        for (int i = 0; i < N; i++) {
            os << "\n[";
            for (int j = 0; j < N; j++) os << p[i][j] << (j + 1 < N ? "," : "");
            os << "]";
        }
        os << "]\n";
        return os;
    }

  private:
    std::pair<int, T> gauss_jordan() {
        int rank = 0;
        T det = 1;
        for (int j = 0; j < N; j++) {
            int pivot = -1;
            for (int i = rank; i < N; i++) {
                if ((*this)[i][j] != T(0)) {
                    pivot = i;
                    break;
                }
            }
            if (pivot == -1) {
                det = 0;
                continue;
            }
            if (pivot != rank) {
                det = -det;
                std::swap((*this)[pivot], (*this)[rank]);
            }
            det *= A[rank][j];
            if (A[rank][j] != T(1)) {
                T coef = T(1) / (*this)[rank][j];
                for (int k = j; k < N; k++) (*this)[rank][k] *= coef;
            }
            for (int i = 0; i < N; i++) {
                if (i == rank) continue;
                T coef = (*this)[i][j];
                if (coef == T(0)) continue;
                for (int k = j; k < N; k++) (*this)[i][k] -= (*this)[rank][k] * coef;
            }
            rank++;
        }
        return {rank, det};
    }
};

#line 2 "src/matrix/SquareMatrix.hpp"
#include <array>
#include <cassert>
#include <utility>

template <typename T, int N> struct SquareMatrix {
    std::array<std::array<T, N>, N> A;

    SquareMatrix() : A() {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                A[i][j] = T(0);
            }
        }
    }

    int size() const { return N; }

    inline const std::array<T, N>& operator[](int i) const { return A[i]; }

    inline std::array<T, N>& operator[](int i) { return A[i]; }

    static SquareMatrix identity() {
        SquareMatrix res;
        for (int i = 0; i < N; i++) res[i][i] = 1;
        return res;
    }

    SquareMatrix& operator+=(const SquareMatrix& B) {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] += B[i][j];
            }
        }
        return *this;
    }

    SquareMatrix& operator-=(const SquareMatrix& B) {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] -= B[i][j];
            }
        }
        return *this;
    }

    SquareMatrix& operator*=(const SquareMatrix& B) {
        std::array<std::array<T, N>, N> C;
        for (int i = 0; i < N; i++) {
            for (int k = 0; k < N; k++) {
                for (int j = 0; j < N; j++) {
                    C[i][j] += (*this)[i][k] * B[k][j];
                }
            }
        }
        std::swap(A, C);
        return *this;
    }

    SquareMatrix& operator*=(const T& v) {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] *= v;
            }
        }
        return *this;
    }

    SquareMatrix& operator/=(const T& v) {
        T inv = T(1) / v;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                (*this)[i][j] *= inv;
            }
        }
        return *this;
    }

    SquareMatrix operator-() const {
        SquareMatrix res;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                res[i][j] = -(*this)[i][j];
            }
        }
        return res;
    }

    SquareMatrix operator+(const SquareMatrix& B) const { return SquareMatrix(*this) += B; }

    SquareMatrix operator-(const SquareMatrix& B) const { return SquareMatrix(*this) -= B; }

    SquareMatrix operator*(const SquareMatrix& B) const { return SquareMatrix(*this) *= B; }

    SquareMatrix operator*(const T& v) const { return SquareMatrix(*this) *= v; }

    SquareMatrix operator/(const T& v) const { return SquareMatrix(*this) /= v; }

    bool operator==(const SquareMatrix& B) const { return A == B.A; }

    bool operator!=(const SquareMatrix& B) const { return A != B.A; }

    SquareMatrix pow(long long n) const {
        assert(0 <= n);
        SquareMatrix x = *this, r = identity();
        while (n) {
            if (n & 1) r *= x;
            x *= x;
            n >>= 1;
        }
        return r;
    }

    SquareMatrix transpose() const {
        SquareMatrix res;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                res[j][i] = (*this)[i][j];
            }
        }
        return res;
    }

    int rank() const { return SquareMatrix(*this).gauss_jordan().first; }

    T det() const { return SquareMatrix(*this).gauss_jordan().second; }

    SquareMatrix inv() const {
        SquareMatrix B(*this), C = identity();
        for (int j = 0; j < N; j++) {
            int pivot = -1;
            for (int i = j; i < N; i++) {
                if (B[i][j] != T(0)) {
                    pivot = i;
                    break;
                }
            }
            assert(pivot != -1);
            if (pivot != j) {
                std::swap(B[pivot], B[j]);
                std::swap(C[pivot], C[j]);
            }
            {
                T coef = T(1) / B[j][j];
                for (int k = 0; k < N; k++) {
                    B[j][k] *= coef;
                    C[j][k] *= coef;
                }
            }
            for (int i = 0; i < N; i++) {
                if (i == j) continue;
                T coef = B[i][j];
                if (coef == T(0)) continue;
                for (int k = 0; k < N; k++) {
                    B[i][k] -= B[j][k] * coef;
                    C[i][k] -= C[j][k] * coef;
                }
            }
        }
        return C;
    }

    friend std::ostream& operator<<(std::ostream& os, const SquareMatrix& p) {
        os << "[(" << N << " * " << N << " Matrix)";
        os << "\n[columun sums: ";
        for (int j = 0; j < N; j++) {
            T sum = 0;
            for (int i = 0; i < N; i++) sum += p[i][j];
            ;
            os << sum << (j + 1 < N ? "," : "");
        }
        os << "]";
        for (int i = 0; i < N; i++) {
            os << "\n[";
            for (int j = 0; j < N; j++) os << p[i][j] << (j + 1 < N ? "," : "");
            os << "]";
        }
        os << "]\n";
        return os;
    }

  private:
    std::pair<int, T> gauss_jordan() {
        int rank = 0;
        T det = 1;
        for (int j = 0; j < N; j++) {
            int pivot = -1;
            for (int i = rank; i < N; i++) {
                if ((*this)[i][j] != T(0)) {
                    pivot = i;
                    break;
                }
            }
            if (pivot == -1) {
                det = 0;
                continue;
            }
            if (pivot != rank) {
                det = -det;
                std::swap((*this)[pivot], (*this)[rank]);
            }
            det *= A[rank][j];
            if (A[rank][j] != T(1)) {
                T coef = T(1) / (*this)[rank][j];
                for (int k = j; k < N; k++) (*this)[rank][k] *= coef;
            }
            for (int i = 0; i < N; i++) {
                if (i == rank) continue;
                T coef = (*this)[i][j];
                if (coef == T(0)) continue;
                for (int k = j; k < N; k++) (*this)[i][k] -= (*this)[rank][k] * coef;
            }
            rank++;
        }
        return {rank, det};
    }
};

Back to top page