#408291 (C++17) No.950 行列累乗

提出ソース

結果

問題	No.950 行列累乗
ユーザー	👑 emthrm
提出日時	2019-12-13 21:45:39
言語	C++17 (gcc 12.3.0 + boost 1.83.0)
結果	WA
実行時間	-
コード長	13,616 bytes
コンパイル時間	1,784 ms
コンパイル使用メモリ	143,880 KB
実行使用メモリ	4,500 KB
最終ジャッジ日時	2023-09-09 23:10:10
合計ジャッジ時間	7,007 ms
ジャッジサーバーID （参考情報）	judge13 / judge12

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	WA	-
testcase_01	AC	2 ms 4,376 KB
testcase_02	WA	-
testcase_03	AC	2 ms 4,380 KB
testcase_04	RE	-
testcase_05	WA	-
testcase_06	AC	2 ms 4,380 KB
testcase_07	WA	-
testcase_08	WA	-
testcase_09	AC	1 ms 4,376 KB
testcase_10	WA	-
testcase_11	WA	-
testcase_12	AC	2 ms 4,376 KB
testcase_13	RE	-
testcase_14	RE	-
testcase_15	RE	-
testcase_16	RE	-
testcase_17	WA	-
testcase_18	AC	2 ms 4,380 KB
testcase_19	WA	-
testcase_20	WA	-
testcase_21	AC	1 ms 4,376 KB
testcase_22	WA	-
testcase_23	WA	-
testcase_24	AC	2 ms 4,376 KB
testcase_25	WA	-
testcase_26	WA	-
testcase_27	AC	2 ms 4,376 KB
testcase_28	WA	-
testcase_29	WA	-
testcase_30	AC	2 ms 4,376 KB
testcase_31	RE	-
testcase_32	RE	-
testcase_33	RE	-
testcase_34	RE	-
testcase_35	RE	-
testcase_36	AC	1 ms 4,376 KB
testcase_37	WA	-
testcase_38	WA	-
testcase_39	AC	2 ms 4,380 KB
testcase_40	WA	-
testcase_41	WA	-
testcase_42	AC	2 ms 4,376 KB
testcase_43	WA	-
testcase_44	WA	-
testcase_45	AC	1 ms 4,376 KB
testcase_46	WA	-
testcase_47	RE	-
testcase_48	WA	-
testcase_49	WA	-
testcase_50	RE	-
testcase_51	RE	-
testcase_52	WA	-
testcase_53	WA	-
testcase_54	RE	-
testcase_55	RE	-
testcase_56	WA	-
testcase_57	WA	-
testcase_58	RE	-
testcase_59	AC	2 ms 4,376 KB
testcase_60	AC	2 ms 4,380 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cctype>
#include <chrono>
#define _USE_MATH_DEFINES
#include <cmath>
#include <cstring>
#include <ctime>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
using namespace std;

#define FOR(i,m,n) for(int i=(m);i<(n);++i)
#define REP(i,n) FOR(i,0,n)
#define ALL(v) (v).begin(),(v).end()

const int INF = 0x3f3f3f3f;
const long long LINF = 0x3f3f3f3f3f3f3f3fLL;
const double EPS = 1e-8;
const int MOD = 1000000007;
// const int MOD = 998244353;
const int dy[] = {1, 0, -1, 0}, dx[] = {0, -1, 0, 1};
// const int dy[] = {1, 1, 0, -1, -1, -1, 0, 1},
//           dx[] = {0, -1, -1, -1, 0, 1, 1, 1};

struct IOSetup {
  IOSetup() {
    cin.tie(nullptr);
    ios_base::sync_with_stdio(false);
    cout << fixed << setprecision(20);
    cerr << fixed << setprecision(10);
  }
} iosetup;
/*-------------------------------------------------*/
int mod = MOD;
struct ModInt {
  unsigned val;
  ModInt(): val(0) {}
  ModInt(long long x) : val(x >= 0 ? x % mod : x % mod + mod) {}
  ModInt pow(long long exponent) {
    ModInt tmp = *this, res = 1;
    while (exponent > 0) {
      if (exponent & 1) res *= tmp;
      tmp *= tmp;
      exponent >>= 1;
    }
    return res;
  }
  ModInt &operator+=(const ModInt &rhs) { if((val += rhs.val) >= mod) val -= mod; return *this; }
  ModInt &operator-=(const ModInt &rhs) { if((val += mod - rhs.val) >= mod) val -= mod; return *this; }
  ModInt &operator*=(const ModInt &rhs) { val = static_cast<unsigned long long>(val) * rhs.val % mod; return *this; }
  ModInt &operator/=(const ModInt &rhs) { return *this *= rhs.inv(); }
  bool operator==(const ModInt &rhs) const { return val == rhs.val; }
  bool operator!=(const ModInt &rhs) const { return val != rhs.val; }
  bool operator<(const ModInt &rhs) const { return val < rhs.val; }
  bool operator<=(const ModInt &rhs) const { return val <= rhs.val; }
  bool operator>(const ModInt &rhs) const { return val > rhs.val; }
  bool operator>=(const ModInt &rhs) const { return val >= rhs.val; }
  ModInt &operator++() { if (++val == mod) val = 0; return *this; }
  ModInt operator++(int) { ModInt res = *this; ++*this; return res; }
  ModInt &operator--() { val = (val == 0 ? mod : val) - 1; return *this; }
  ModInt operator--(int) { ModInt res = *this; --*this; return res; }
  ModInt operator+() const { return *this; }
  ModInt operator-() const { return ModInt(val ? mod - val : 0); }
  ModInt operator+(const ModInt &rhs) const { return ModInt(*this) += rhs; }
  ModInt operator-(const ModInt &rhs) const { return ModInt(*this) -= rhs; }
  ModInt operator*(const ModInt &rhs) const { return ModInt(*this) *= rhs; }
  ModInt operator/(const ModInt &rhs) const { return ModInt(*this) /= rhs; }
  friend ostream &operator<<(ostream &os, const ModInt &rhs) { return os << rhs.val; }
  friend istream &operator>>(istream &is, ModInt &rhs) { long long x; is >> x; rhs = ModInt(x); return is; }
private:
  ModInt inv() const {
    // if (__gcd(val, mod) != 1) assert(false);
    unsigned a = val, b = mod; int x = 1, y = 0;
    while (b) {
      unsigned tmp = a / b;
      swap(a -= tmp * b, b);
      swap(x -= tmp * y, y);
    }
    return ModInt(x);
  }
};
ModInt abs(const ModInt &x) { return x; }
struct Combinatorics {
  int val;
  vector<ModInt> fact, fact_inv, inv;
  Combinatorics(int val = 10000000) : val(val), fact(val + 1), fact_inv(val + 1), inv(val + 1) {
    fact[0] = 1;
    FOR(i, 1, val + 1) fact[i] = fact[i - 1] * i;
    fact_inv[val] = ModInt(1) / fact[val];
    for (int i = val; i > 0; --i) fact_inv[i - 1] = fact_inv[i] * i;
    FOR(i, 1, val + 1) inv[i] = fact[i - 1] * fact_inv[i];
  }
  ModInt nCk(int n, int k) {
    if (n < 0 || n < k || k < 0) return ModInt(0);
    // assert(n <= val && k <= val);
    return fact[n] * fact_inv[k] * fact_inv[n - k];
  }
  ModInt nPk(int n, int k) {
    if (n < 0 || n < k || k < 0) return ModInt(0);
    // assert(n <= val);
    return fact[n] * fact_inv[n - k];
  }
  ModInt nHk(int n, int k) {
    if (n < 0 || k < 0) return ModInt(0);
    return (k == 0 ? ModInt(1) : nCk(n + k - 1, k));
  }
};

template <typename T>
struct Matrix {
  Matrix(int m, int n, T val = 0) : dat(m, vector<T>(n, val)) {}

  int height() const { return dat.size(); }

  int width() const { return dat.front().size(); }

  Matrix pow(long long exponent) {
    int n = height();
    Matrix<T> tmp = *this, res(n, n, 0);
    REP(i, n) res[i][i] = 1;
    while (exponent > 0) {
      if (exponent & 1) res *= tmp;
      tmp *= tmp;
      exponent >>= 1;
    }
    return res;
  }

  inline const vector<T> &operator[](const int idx) const { return dat[idx]; }
  inline vector<T> &operator[](const int idx) { return dat[idx]; }

  Matrix &operator=(const Matrix &rhs) {
    int m = rhs.height(), n = rhs.width();
    dat.resize(m, vector<T>(n));
    REP(i, m) REP(j, n) dat[i][j] = rhs[i][j];
    return *this;
  }

  Matrix &operator+=(const Matrix &rhs) {
    int m = height(), n = width();
    REP(i, m) REP(j, n) dat[i][j] += rhs[i][j];
    return *this;
  }

  Matrix &operator-=(const Matrix &rhs) {
    int m = height(), n = width();
    REP(i, m) REP(j, n) dat[i][j] -= rhs[i][j];
    return *this;
  }

  Matrix &operator*=(const Matrix &rhs) {
    int m = height(), n = rhs.width(), l = width();
    vector<vector<T> > res(m, vector<T>(n, 0));
    REP(i, m) REP(j, n) {
      REP(k, l) res[i][j] += dat[i][k] * rhs[k][j];
    }
    swap(dat, res);
    return *this;
  }

  Matrix operator+(const Matrix &rhs) const { return Matrix(*this) += rhs; }

  Matrix operator-(const Matrix &rhs) const { return Matrix(*this) -= rhs; }

  Matrix operator*(const Matrix &rhs) const { return Matrix(*this) *= rhs; }

private:
  vector<vector<T> > dat;
};

template <typename T>
int gauss_jordan(Matrix<T> &mat, bool is_extended = false) {
  int m = mat.height(), n = mat.width(), rank = 0;
  REP(col, n) {
    if (is_extended && col == n - 1) break;
    int pivot = -1;
    T mx = EPS;
    FOR(row, rank, m) {
      if (abs(mat[row][col]) > mx) {
        pivot = row;
        mx = abs(mat[row][col]);
      }
    }
    if (pivot == -1) continue;
    swap(mat[rank], mat[pivot]);
    T tmp = mat[rank][col];
    REP(col2, n) mat[rank][col2] /= tmp;
    REP(row, m) if (row != rank && abs(mat[row][col]) > EPS) {
      tmp = mat[row][col];
      REP(col2, n) mat[row][col2] -= mat[rank][col2] * tmp;
    }
    ++rank;
  }
  return rank;
}

template <typename T, typename U = double>
vector<U> linear_equation(const Matrix<T> &lhs, const vector<T> &rhs) {
  int m = lhs.height(), n = lhs.width();
  Matrix<U> matrix(m, n + 1);
  REP(i, m) {
    REP(j, n) matrix[i][j] = lhs[i][j];
    matrix[i][n] = rhs[i];
  }
  int rank = gauss_jordan(matrix, true);
  vector<U> res;
  FOR(row, rank, m) {
    if (abs(matrix[row][n]) > EPS) return res;
  }
  res.assign(n, 0);
  REP(i, rank) res[i] = matrix[i][n];
  return res;
}

template <typename T, typename U = double>
bool inverse(const Matrix<T> &mat, Matrix<U> &inv) {
  int n = mat.height();
  Matrix<U> gauss_jordan(n, n << 1, U(0));
  REP(i, n) {
    REP(j, n) gauss_jordan[i][j] = mat[i][j];
    gauss_jordan[i][n + i] = U(1);
  }
  REP(col, n) {
    int pivot = -1;
    U mx = EPS;
    FOR(row, col, n) {
      if (abs(gauss_jordan[row][col]) > mx) {
        pivot = row;
        mx = abs(gauss_jordan[row][col]);
      }
    }
    if (pivot == -1) return false;
    swap(gauss_jordan[col], gauss_jordan[pivot]);
    U tmp = gauss_jordan[col][col];
    REP(col2, n << 1) gauss_jordan[col][col2] /= tmp;
    REP(row, n) if (row != col && abs(gauss_jordan[row][col]) > EPS) {
      tmp = gauss_jordan[row][col];
      REP(col2, n << 1) gauss_jordan[row][col2] -= gauss_jordan[col][col2] * tmp;
    }
  }
  assert(inv.height() == n && inv.width() == n);
  REP(i, n) REP(j, n) inv[i][j] = gauss_jordan[i][n + j];
  return true;
}

long long mod_pow(long long base, long long exponent, int mod = MOD) {
  base %= mod;
  long long res = 1;
  while (exponent > 0) {
    if (exponent & 1) (res *= base) %= mod;
    (base *= base) %= mod;
    exponent >>= 1;
  }
  return res;
}

long long mod_inv(long long a, int mod = MOD) { return mod_pow(a % mod, mod - 2, mod); }

long long bsgs(long long g, long long y, int mod = MOD) {
  g %= mod;
  if (g < 0) g += mod;
  y %= mod;
  if (y < 0) y += mod;
  int root = ceil(sqrt(mod));
  map<long long, int> g_to_the_jth_power;
  long long tmp = 1 % mod;
  REP(i, root) {
    if (g_to_the_jth_power.count(tmp) == 0) g_to_the_jth_power[tmp] = i;
    (tmp *= g) %= mod;
  }
  long long inv = mod_pow(mod_inv(g, mod), root, mod), rhs = y;
  REP(i, root) {
    if (g_to_the_jth_power.count(rhs) == 1) return static_cast<long long>(i) * root + g_to_the_jth_power[rhs];
    (rhs *= inv) %= mod;
  }
  return -1;
}

long long bsgs_not0(long long g, long long y, int mod = MOD) {
  g %= mod;
  if (g < 0) g += mod;
  y %= mod;
  if (y < 0) y += mod;
  int root = ceil(sqrt(mod));
  map<long long, int> g_to_the_jth_power;
  long long tmp = 1 % mod;
  REP(i, root) {
    if (g_to_the_jth_power.count(tmp) == 0) g_to_the_jth_power[tmp] = i;
    (tmp *= g) %= mod;
  }
  long long inv = mod_pow(mod_inv(g, mod), root, mod), rhs = y;
  REP(i, root) {
    if (g_to_the_jth_power.count(rhs) == 1) {
      long long res = static_cast<long long>(i) * root + g_to_the_jth_power[rhs];
      if (res > 0) return res;
    }
    (rhs *= inv) %= mod;
  }
  return -1;
}

pair<long long, long long> ext_gcd(long long a, long long b) {
  if (b == 0) return {1, 0};
  pair<long long, long long> pr = ext_gcd(b, a % b);
  return {pr.second, pr.first - a / b * pr.second};
}

// https://judge.yosupo.jp/submission/318
template< typename T >
T mod_pow(T x, T n, const T &p) {
  T ret = 1;
  while(n > 0) {
    if(n & 1) (ret *= x) %= p;
    (x *= x) %= p;
    n >>= 1;
  }
  return ret;
}
template< typename T >
T mod_sqrt(const T &a, const T &p) {
  if(a == 0) return 0;
  if(p == 2) return a;
  if(mod_pow(a, (p - 1) >> 1, p) != 1) {
    cout << -1 << '\n';
    exit(0);
  }
  T b = 1;
  while(mod_pow(b, (p - 1) >> 1, p) == 1) ++b;
  T e = 0, m = p - 1;
  while(m % 2 == 0) m >>= 1, ++e;
  T x = mod_pow(a, (m - 1) >> 1, p);
  T y = a * (x * x % p) % p;
  (x *= a) %= p;
  T z = mod_pow(b, m, p);
  while(y != 1) {
    T j = 0, t = y;
    while(t != 1) {
      j += 1;
      (t *= t) %= p;
    }
    z = mod_pow(z, T(1) << (e - j - 1), p);
    (x *= z) %= p;
    (z *= z) %= p;
    (y *= z) %= p;
    e = j;
  }
  return x;
}

int main() {
  cin >> mod;
  Matrix<ModInt> a(2, 2), b(2, 2);
  REP(i, 2) REP(j, 2) cin >> a[i][j];
  REP(i, 2) REP(j, 2) cin >> b[i][j];
  long long D = ((a[0][0] + a[1][1]).pow(2) - (a[0][0] * a[1][1] - a[0][1] * a[1][0]) * 4).val, sqrtD = mod_sqrt(D, 1LL * mod);
  ModInt lambda1 = (a[0][0] + a[1][1] + D) / 2, lambda2 = (a[0][0] + a[1][1] - D) / 2;
  if (lambda1 == lambda2) {
    // Jordan normal form
    assert(false);
  } else {
    Matrix<ModInt> m(2, 2);
    vector<ModInt> ZERO(2, 0);
    m[0][0] = a[0][0] - lambda1; m[0][1] = a[0][1];
    m[1][0] = a[1][0];           m[1][1] = a[1][1] - lambda1;
    vector<ModInt> P_l = linear_equation<ModInt, ModInt>(m, ZERO);
    if (P_l.empty()) {
      cout << -1 << '\n';
      return 0;
    }
    m[0][0] = a[0][0] - lambda2; m[0][1] = a[0][1];
    m[1][0] = a[1][0];           m[1][1] = a[1][1] - lambda2;
    vector<ModInt> P_r = linear_equation<ModInt, ModInt>(m, ZERO);
    if (P_r.empty()) {
      cout << -1 << '\n';
      return 0;
    }
    Matrix<ModInt> P(2, 2), inv(2, 2);
    P[0][0] = P_l[0]; P[0][1] = P_r[0];
    P[1][0] = P_l[1]; P[1][1] = P_r[1];
    if (!inverse(P, inv)) {
      cout << -1 << '\n';
      return 0;
    }
    Matrix<ModInt> lhs(4, 2);
    vector<ModInt> rhs(4, 1);
    lhs[0][0] = P[0][0] * inv[0][0]; lhs[0][1] = P[0][1] * inv[1][0];
    lhs[1][0] = P[0][0] * inv[1][0]; lhs[1][1] = P[0][1] * inv[1][1];
    lhs[2][0] = P[1][0] * inv[0][0]; lhs[2][1] = P[1][1] * inv[1][0];
    lhs[3][0] = P[1][0] * inv[1][0]; lhs[3][1] = P[1][1] * inv[1][1];
    rhs[0] = b[0][0];
    rhs[1] = b[0][1];
    rhs[2] = b[1][0];
    rhs[3] = b[1][1];
    vector<ModInt> pow_n = linear_equation<ModInt, ModInt>(lhs, rhs);
    if (pow_n.empty()) {
      cout << -1 << '\n';
      return 0;
    }
    long long n1 = bsgs_not0(lambda1.val, pow_n[0].val), n2 = bsgs_not0(lambda2.val, pow_n[1].val);
    long long ans = LINF;
    if (n1 != -1 && lambda2.pow(n1) == pow_n[1]) ans = n1;
    if (n2 != -1 && lambda1.pow(n2) == pow_n[0]) ans = min(ans, n2);
    if (ans != LINF) {
      cout << ans << '\n';
      return 0;
    }
    n1 = bsgs(lambda1.val, pow_n[0].val); n2 = bsgs(lambda2.val, pow_n[1].val);
    long long ONE1 = bsgs_not0(lambda1.val, 1), ONE2 = bsgs_not0(lambda2.val, 1);
    // n1 + ONE1 * k = n2 + ONE2 * l を満たす自然数 k, l を求める
    // ONE1 * k - ONE2 * l = n2-n1
    if (n1 > n2) {
      swap(n1, n2);
      swap(ONE1, ONE2);
    } else if (n1 == n2) {
      cout << n1 << '\n';
      return 0;
    }
    if ((n2 - n1) % __gcd(ONE1, ONE2) != 0) {
      cout << -1 << '\n';
      return 0;
    }
    auto pr = ext_gcd(ONE1, ONE2);
    if (n1 + pr.first * ONE1 > 0) {
      cout << (n2 - n1) / __gcd(ONE1, ONE2) * (n1 + pr.first * ONE1) << '\n';
    } else if (n2 + pr.second * ONE2 > 0) {
      cout << (n2 - n1) / __gcd(ONE1, ONE2) * (n2 + pr.second * ONE2) << '\n';
    } else {
      assert(false);
    }
  }
  return 0;
}

yukicoder

結果

テストケース

ソースコード