linalg/diagonalize_triangular_matrix.hpp

• View this file on GitHub
• Last update: 2025-07-02 18:15:14+09:00
• Include: #include "linalg/diagonalize_triangular_matrix.hpp"
• Link: https://codeforces.com/contest/1540/problem/E

Depends on

• linalg/matrix_inv.hpp
• linalg/transpose.hpp

Code

#include "linalg/transpose.hpp"
#include "linalg/matrix_inv.hpp"

// P^{-1}AP = diag(D) となる (D,P,P^{-1}) を返す
// https://codeforces.com/contest/1540/problem/E
template <typename mint>
tuple<vc<mint>, vvc<mint>, vvc<mint>> diagonalize_triangular_matrix(
    vvc<mint> A) {
  int N = len(A);
  bool lower = 1, upper = 1;
  FOR(i, N) FOR(j, N) {
    if (i < j) lower = lower && (A[i][j] == 0);
    if (i > j) upper = upper && (A[i][j] == 0);
  }
  assert(lower || upper);
  if (!upper) {
    A = transpose(A);
    auto [D, P, Q] = diagonalize_triangular_matrix(A);
    return {D, transpose(Q), transpose(P)};
  }
  assert(upper);

  vc<mint> D(N);
  vv(mint, P, N, N);
  FOR(k, N) D[k] = A[k][k];
  FOR(j, N) FOR(i, j) assert(D[i] != D[j]);

  FOR(k, N) {
    mint d = A[k][k];
    D[k] = d;
    vc<mint> X(N);
    X[k] = 1;
    FOR_R(i, k) {
      mint x = 0;
      FOR(j, i + 1, k + 1) x += A[i][j] * X[j];
      X[i] = x / (d - D[i]);
    }
    FOR(i, N) P[i][k] = X[i];
  }
  auto Q = matrix_inv<mint>(P).se;
  return {D, P, Q};
}

#line 1 "linalg/diagonalize_triangular_matrix.hpp"

#line 1 "linalg/transpose.hpp"
template <typename VC>
vc<VC> transpose(const vc<VC>& A, int H = -1, int W = -1) {
  if (H == -1) { H = len(A), W = (len(A) == 0 ? 0 : len(A[0])); }
  vc<VC> B(W, VC(H, 0));
  FOR(x, H) FOR(y, W) B[y][x] = A[x][y];
  return B;
}
#line 2 "linalg/matrix_inv.hpp"

// (det, invA) をかえす

template <typename T>
pair<T, vc<vc<T>>> matrix_inv(vc<vc<T>> A) {
  T det = 1;
  int N = len(A);
  vv(T, B, N, N);
  FOR(n, N) B[n][n] = 1;
  FOR(i, N) {
    FOR(k, i, N) if (A[k][i] != 0) {
      if (k != i) {
        swap(A[i], A[k]), swap(B[i], B[k]);
        det = -det;
      }
      break;
    }
    if (A[i][i] == 0) return {T(0), {}};
    T c = T(1) / A[i][i];
    det *= A[i][i];
    FOR(j, i, N) A[i][j] *= c;
    FOR(j, N) B[i][j] *= c;
    FOR(k, N) if (i != k) {
      T c = A[k][i];
      FOR(j, i, N) A[k][j] -= A[i][j] * c;
      FOR(j, N) B[k][j] -= B[i][j] * c;
    }
  }
  return {det, B};
}
#line 4 "linalg/diagonalize_triangular_matrix.hpp"

// P^{-1}AP = diag(D) となる (D,P,P^{-1}) を返す
// https://codeforces.com/contest/1540/problem/E
template <typename mint>
tuple<vc<mint>, vvc<mint>, vvc<mint>> diagonalize_triangular_matrix(
    vvc<mint> A) {
  int N = len(A);
  bool lower = 1, upper = 1;
  FOR(i, N) FOR(j, N) {
    if (i < j) lower = lower && (A[i][j] == 0);
    if (i > j) upper = upper && (A[i][j] == 0);
  }
  assert(lower || upper);
  if (!upper) {
    A = transpose(A);
    auto [D, P, Q] = diagonalize_triangular_matrix(A);
    return {D, transpose(Q), transpose(P)};
  }
  assert(upper);

  vc<mint> D(N);
  vv(mint, P, N, N);
  FOR(k, N) D[k] = A[k][k];
  FOR(j, N) FOR(i, j) assert(D[i] != D[j]);

  FOR(k, N) {
    mint d = A[k][k];
    D[k] = d;
    vc<mint> X(N);
    X[k] = 1;
    FOR_R(i, k) {
      mint x = 0;
      FOR(j, i + 1, k + 1) x += A[i][j] * X[j];
      X[i] = x / (d - D[i]);
    }
    FOR(i, N) P[i][k] = X[i];
  }
  auto Q = matrix_inv<mint>(P).se;
  return {D, P, Q};
}

Back to top page