Baby Stepsなブログ

競プロとか。間違ったこと書いてあったら@pi0nep1oneにご連絡ください。

行列累乗についてのメモと出題例

最近DPの漸化式を行列式に変換するセンスを磨くために関連問題を解いており、問題や行列累乗全般について雑多なメモを残しておこうと思う。

フィボナッチ数列のN項目を高速に求める

行列累乗の有名な活用例としてよくフィボナッチ数列が挙げられる。

オンラインジャッジで全く同じ問題は見つけられなかったが、蟻本にその記載がある。

フィボナッチ数列の漸化式は以下の様になる。

  • dp[0]=dp[1]=1
  • dp[i+2]=dp[i+1]+dp[i]

この遷移をナイーブに実装すると第N項を求めるためにO(N)の時間が掛かる。

これを高速に求めるため、上記の漸化式を行列積の形に書き換える。(f_iをフィボナッチ数列の第i項目とすると)

{f_1, f_0} = {{1,1}, {1,0}} * {1,0}

更に右辺に {{1,1}, {1,0}} を掛けると {f_2, f_1}、更に掛けると {f_3, f_2} を求めることができる。

つまり第N項目を求めるためには {{1,1}, {1,0}} のN乗を求めればよい。累乗を高速に求める手法に繰り返し二乗法というものがあるが、これが行列にも応用することができる。この高速に行列の累乗を求める手法を競プロ界隈では「行列累乗」と呼称しているらしい。

以下は、行列累乗によりフィボナッチ数列を求めるための実装

フィボナッチ数列

#define mat vector<vector<ll>>

/// 行列積
mat mat_mul(mat &a, mat &b) {
  mat res(a.size(), vector<ll>(b[0].size()));
  for (int i = 0; i < a.size(); i++) {
    for (int j = 0; j < b[0].size(); j++) {
      for (int k = 0; k < b.size(); k++) {
        (res[i][j] += a[i][k] * b[k][j]) %= MOD;
      }
    }
  }
  return res;
}

/// 行列累乗
mat mat_pow(mat a, long long n) {
  mat res(a.size(), vector<ll>(a.size()));
  // 単位行列で初期化
  for (int i = 0; i < a.size(); i++)
    res[i][i] = 1;
  // 繰り返し二乗法
  while (n > 0) {
    if (n & 1) res = mat_mul(a, res);
    a = mat_mul(a, a);
    n >>= 1;
  }
  return res;
}

signed main()
{
  ll n;
  cin >> n;
  
  // フィボナッチ数列の行列累乗
  mat m(2, vc<ll>(2));
  m[0][0] = m[0][1] = m[1][0] = 1;
  m[1][1] = 0;
  m = mat_pow(m, n);
  ll ans = m[1][0];
  cout << ans << endl;
}

例題をいくつか

以下に行列累乗によって解ける問題を載せておく。

下に行くほど難しくなっている。

ABC 009 D - 漸化式

atcoder.jp

この問題で漸化式を立てると以下の様になる。

A[i] = (C[0]&A[i-1]) ^ (C[1]&A[i-2]) ^ ...

この遷移を観察すると、自身より前の項によって自身の項が定まるためフィボナッチ数列DPに似ている。

また、求めたい第M項のMの値が大きい(最大109)かつ、自身を求めるために参照する自身より前の項の数Kが小さい(最大100)であるということから計算量から行列累乗が解法の候補に挙がる。(行列累乗で解けるなら、Kが行列の列幅行幅に対応するので、その計算量はO(K3logM)と十分に高速である。)

ここで、この問題では演算が加法、乗法ではなく、排他的論理和(XOR)、論理積(AND)を使うが半環を満たすため行列累乗が使える。

累乗する行列は以下の様になる。

※ 以下の画像はは解説スライドのP.102の引用です。

この問題では乗法が論理積に置き換わっているので、ここでの1とは、論理積を行った結果が変化しない値である。Aの要素は32ビット符号なし整数なので、32ビット1が立った値、(1<<32)-1 などを扱えばよい。

f:id:pione8888:20210404202641j:plain

解説スライド

www.slideshare.net

解答コード

atcoder.jp

ARC 050 C - LCM 111

atcoder.jp

X桁のレピュニット数をf(X)と表すことにする。

レピュニット数同士の剰余を観察してみると(A>Bとする)

f(A)%f(B) = f(A%B)

という関係があることがわかる。レピュニット数同士の剰余がその桁数の剰余に対応することから、レピュニット数同士のgcdはその桁数のgcdによって求めることができる。つまり、

gcd(f(A), f(B)) = f(gcd(A, B))

という関係が成り立つ。

gcd(A, B)をGとすると求める解は、

f(A)*f(B)/f(G)

で求めることができるので、後は行列累乗でf(A), f(B), f(G)を高速に求めれば良さそう...に思えるが、この問題で与えられる法Mは適当な正整数であり(f(G)と互いに素であるとは限らないため)f(G)の逆元を求められない可能性がある。

ひとまず f(A) を求めることを考える。

f(i+1) = f(i) * 10 + 1

なので、これを行列式で表すと以下の様になる。

{f(2), f(1)} = {{10,1},{0,1}}*{1,1}

よって、f(N)を求めるには、{{10,1},{0,1}}のN乗を行列累乗で求めればよいとわかる。

残る f(B)/f(G) について、これは逆元を使わずに求めることにする。f(B), f(G)はいずれもレピュニット数で f(B)>=f(G) なので、このような関係性の2数の除法について観察してみる。

f(6) / f(2) = 111111 / 11 = 10101

この 10101 という値はレピュニット数の遷移とほぼ同じで求めることができる。つまり、以下の様に表せる。

g(i+1) = g(i) * 10G + 1

これを行列式で表すと以下の様になる。

{g(2), g(1)} = {{10G,1},{0,1}}*{1,1}

B/G = g とすると、f(B)/f(G) は {{10G,1},{0,1}}*{1,1} をg乗することで求めることができる。

解答コード

atcoder.jp

EDPC R - Walk

atcoder.jp

この問題もKという値が大きい(最大1018)かつ、行列の列幅になるであろうNの値が小さい(最大50)ので、行列累乗で解けるならO(N3logK)時間になり間に合いそう。(さらにこの問題は数え上げである)

まず素直に漸化式を考えると、以下の様な遷移が考えられる。

dp[len][pos] := 経路長lenで終点がposであるものの通り数

  • dp[0][i] = 1 (初期化)
  • dp[len+1][to] += dp[len][from]

簡単のため N=2 という小さい場合について考える。有向辺は以下の様に与えられるとする。

0 0
1 0 // 頂点2->頂点1の有向辺

この行列を(i, j)=(頂点i <- 頂点j の有向辺が存在する)というように変換する。

0 1 // 頂点1<-頂点2の有向辺
0 0

この行列を使えば、以下の様に行列式で漸化式を表すことができる。

{dp[1][0], dp[1][1]} = {{0,1}, {0,0}} * {1,1}

経路長2の長さを求めるに右辺に {{0,1}, {0,0}} を掛ければよく、つまり経路長Kの長さを求めるにはこのK乗を求めればよい。

以上を有向グラフ全般に適用させるためには、与えられる行列Aから、(頂点i <- 頂点j の有向辺が存在する)行列を構築してそのK乗を求めればよい。

解答コード

atcoder.jp

yukicoder No.840 ほむほむほむら

yukicoder.me

h(ほ) -> m(む) -> r(ら) の並び順の部分文字列の個数がK倍個であるような文字列の個数を数え上げる問題。

この問題も、文字列の長さとして与えられるNが大きく(最大109)、Kが小さい(最大5)。

まず素直にDPの漸化式を考えてみる。

dp[n][i][j][k] := n文字で、hの個数がi個の倍数、hmの個数がj個の倍数、hmrの個数がkの倍数である場合の通り数(よって配列長はN*K3

  • dp[0][0][0][0] = 1 (初期化)
  • dp[n][(i+1)%K][j][k] += dp[n-1][i][j][k] // hを追加
  • dp[n][i][(j+i)%K][k] += dp[n-1][i][j][k] // mを追加
  • dp[n][i][j][(k+j)%K] += dp[n-1][i][j][k] // rを追加

これを行列式で合わらすために [i][j][k] の3次元配列をK進数によって1次元で表すことを考える。

つまり、hがi個、mがj個、rがk個を iK2+jK+k と表すことにする。この値は高々K3(<=53)と小さく、この行幅列幅の行列を累乗しても十分高速である。

あとは、行列mのm(to, from)を状態toへ状態fromからの遷移があるとして初期化し、このN乗を求めればN文字の各状態の通り数を計算できる。このうち、hmrの個数が0倍の個数の総和を求めればそれが解となる。

解答コード

yukicoder.me