ABC242C別解: 行列による解法

ABC242C の別解

atcoder.jp

解答コード

提出_1: TLE．(N-1)回の行列の内積
提出_2: ギリAC．対称性を利用して行列サイズを小さくする
提出_3: AC．(N-1)回の内積を効率よく計算　　<-本記事はこれ

解説

すでにDP（動的計画法）で求める方法は解説いたしました <リンク>．

今回のDPの計算を行列で表すと，

w = np.array(
    [[1, 1, 0, 0, 0, 0, 0, 0, 0],  # (0, 1) 番目の和
     [1, 1, 1, 0, 0, 0, 0, 0, 0],  # (0, 1, 2) 番目の和
     [0, 1, 1, 1, 0, 0, 0, 0, 0],  # (1, 2, 3) 番目の和
     [0, 0, 1, 1, 1, 0, 0, 0, 0],  # (2, 3, 4) 番目の和
     [0, 0, 0, 1, 1, 1, 0, 0, 0],  # (3, 4, 5) 番目の和
     [0, 0, 0, 0, 1, 1, 1, 0, 0],  # (4, 5, 6) 番目の和
     [0, 0, 0, 0, 0, 1, 1, 1, 0],  # (5, 6, 7) 番目の和
     [0, 0, 0, 0, 0, 0, 1, 1, 1],  # (6, 7, 8) 番目の和
     [0, 0, 0, 0, 0, 0, 0, 1, 1]]) # (7, 8) 番目の和

こんな行列になります．
実行してみると，もちろん前解説のdpを再現しています．

dp = np.array([1]*9)
print(dp)
for i in range(3):
    dp = w@dp
    print(dp)
# [ 1,  1,  1,  1,  1,  1,  1,  1,  1]
# [ 2,  3,  3,  3,  3,  3,  3,  3,  2]
# [ 5,  8,  9,  9,  9,  9,  9,  8,  5]
# [13, 22, 26, 27, 27, 27, 26, 22, 13]

つまりこの問題は，Nが与えられたら，上の行列の(N-1)乗をして，上の行列の要素内の和を求めなさい．というように言い換えることができます．

しかし，行列の計算って遅いんです．

実際にすると提出_1のように，TLEになりました．

そこで，対称性を用いて計算量を減らすと，なんとかACになります（提出_2）．

悩むこと一晩...

そしてついに，よくよく考えたら，行列のN乗を計算するときに，繰り返し自乗法で冪乗の計算の高速化と同様にすれば，早いじゃんっ！って気が付きました．

その関数がこちらです．正則行列をn乗します．途中でmodで割ります．ただそれだけです．

import numpy as np

def pow_array(A, n, mod):
    B = np.eye(len(A), dtype='object')
    while n:
        if n%2:
            B = (B@A) %mod
        A = (A@A) %mod
        n //= 2
    return B

これを書けば，あとは単純です．入力の読み込みとMODを代入，初期の行列を定義して...

MOD = 998244353
N = int(input())
w = np.zeros([9, 9], dtype='int')
for i in range(9):
    for j in (i-1, i, i+1):
        if 0<=j<=8:
            w[i][j] = 1

(N-1)乗を計算して，合計を出力します．

w = pow_array(w, N-1, mod=MOD)
ans = w.sum()%MOD
print(ans)

ライブラリのnumpyのndarrayで計算してますが，
もしかするとnumpyを使用しないで，
内積計算をpythonのリストとforで行えば，
多少処理が速くなるかもしれません．

以上，ABC242Cの別解でした．

＜参考解説リンク＞
・ユーザ解説 by MMNMM