CodeChef October Challenge 2019: Queries on Matrix

CodeChef October Challenge 2019 の問題: Queries on Matrix (Code: JIIT)

Contest Page | CodeChef

問題概要

の行列が与えられる。全ての要素は初め 0 である。

行目列目の要素を として記述する。

以下の操作をちょうど 回行う。

1. 1つの要素 を選ぶ
2. 行目の全ての要素に 1 を足す
3. 列目の全ての要素に 1 を足す

回の操作の後、行列全体の内で要素が奇数となる要素の数を 個とする要素の操作列は何通り存在するかを modulo 998244353 で求めよ。
(操作列は、異なる順序で要素を選ぶもの同士を区別する)

制約

• 1つのテストケースのケース数:
• 
• 
• 

解法

行列累乗で試行錯誤しながら計算した。

まず回の操作の後、行の中で奇数回1を足した行の数を、列の中で奇数回1を足した列の数をとする時、奇数となる要素の数は になるため、これが と等しくなるように を選べばよい。
この選び方は か かで場合分けすればよく、選び方は になる。

行と列については独立に考えられるため、前もって となる行の選び方 と となる列の選び方 をそれぞれ計算し、全ての選び方について計算すればよい。
ここからは行について考える。(列についても同様に求められる)

DPで求める場合、以下で求めることができる。

• = i回操作した後、奇数回選んだ行が 個となる通り数

このDPは以下で計算できる。

• 

最終的に各について を求めればよい。(この計算量は )

次に行列累乗を用いて計算することを考える。
回後の操作において、奇数回選んだ行数が 個となる通り数をとする、長さ のベクトル とし、このベクトルに行列を掛けることで最終的に を求める。
初期値の は となる。

以下を満たす の行列を考える。

• 各 について
  • 
  • 
• それ以外の要素は 0

そして、 を求める。

N = 5 の時の行列

この計算量は となる。

ここから、この行列累乗を高速化する。

行列を対角化する。
まず、行列について固有値を求めてみると、 となることが分かる。

次に各固有値に対する固有ベクトルを求める必要があるが、行列がtridiagonalの形になるため、固有ベクトルを で求めることができる。
具体的には、行列に対するある固有値の1つの固有ベクトル は以下の数列から求まる。

• ,
• 

これらの固有ベクトルから成る行列 を考える。
この時、この行列の逆行列 は になるため、逆行列を改めて求める必要がない。

また、対角行列 の 乗の は で求まる。

よって最終的に求めたいベクトル は で求められる。

N = 5 の例

全体の計算量は

実装

提出コード(C++14): Solution: 27139983 | CodeChef

#define N 2003

const ll mod = 998244353;

int n, m, z;
ll q;

ll xs[N], ys[N];
ll vs[N][N];

ll fast_pow(ll x, ll n) {
  ll res = 1;
  while(n > 0) {
    if(n & 1) {
      res = res * x % mod;
    }
    x = x * x % mod;
    n >>= 1;
  }
  return res;
}

ll revv[N];
void calc(ll xs[N], int n, ll q) {
  rep(i, n+1) xs[i] = 0;
  rep(i, n+1) {
    // 固有ベクトルの計算
    ll x = (mod + n - 2*i) % mod;
    vs[i][0] = 1; vs[i][1] = x;
    for(int j = 2; j <= n; ++j) {
      vs[i][j] = ((x * vs[i][j-1] + vs[i][j-2] * (mod+j-n-2)) % mod) * revv[j] % mod;
    }

    // x_Q の計算
    ll v = fast_pow(x, q % (mod-1)) * vs[0][i] % mod;
    rep(j, n+1) {
      xs[j] += v * vs[i][j];
      xs[j] %= mod;
    }
  }
  ll rev = fast_pow(revv[2], n);
  rep(i, n+1) xs[i] = xs[i] * rev % mod;
}

int main() {
  revv[0] = revv[1] = 1;
  for(int i = 2; i <= 2000; ++i) revv[i] = fast_pow(i, mod-2);
  int t; cin >> t;
  while(t--) {
    cin >> n >> m >> q >> z;
    ll res = 0;
    calc(xs, n, q);
    calc(ys, m, q);

    if(n*m == 2*z) {
      // NM = 2Z の場合は (N - 2a) (M - 2b) = 0 を満たす
      if(n % 2 == 0) {
        rep(b, m+1) {
          if(b * 2 != m) res += xs[n/2] * ys[b] % mod;
        }
      }
      if(m % 2 == 0) {
        rep(a, n+1) {
          if(a * 2 != n) res += xs[a] * ys[m/2] % mod;
        }
      }
      if(n % 2 == 0 && m % 2 == 0) {
        res += xs[n/2] * ys[m/2] % mod;
      }
      res %= mod;
    } else {
      // NM != 2Z の場合は Z = aM + bN - 2ab を満たすa, bを走査
      rep(a, n+1) {
        int p = (z - a*m), q = (n - 2*a);
        if(q < 0) {
          p = -p; q = -q;
        }
        if(q == 0 || p < 0 || p % q != 0) {
          continue;
        }
        int b = p / q;
        if(0 <= b && b <= m) {
          res += xs[a] * ys[b] % mod;
        }
      }
      res %= mod;
    }
    cout << res << "\n";
  }
  return 0;
}