【今日の精進】ABC256G - Black and White Stones

毎日格上の問題を倒すやつの15日目です。

問題リンク

atcoder.jp

問題概要

一辺が の正 角形に白黒の石を置くとき、すべての辺の黒い石の数が同じ置き方の数を で割ったあまりを求めてね。

解法

一辺に置く黒石の数を決め打つ余裕はありますが、愚直なDPで計算する余裕はありません。¹

愚直なDPでは、角にある白と黒の石で場合分けして二項係数を使います。

また、遷移が変わらないため、行列累乗を使うことができます。

最後は元の角に戻るため、「白→白」または「黒→黒」のみをカウントしないといけません。

提出コード

def power(mat,num):
    if num==1:
        return mat
    bef=power(mat,num//2)
    aft=[[0,0],[0,0]]
    for i in range(2):
        for j in range(2):
            for k in range(2):
                aft[i][k]+=bef[i][j]*bef[j][k]
                aft[i][k]%=MOD
    if num%2==0:
        return aft
    else:
        ret=[[0,0],[0,0]]
        for i in range(2):
            for j in range(2):
                for k in range(2):
                    ret[i][k]+=aft[i][j]*mat[j][k]
                    ret[i][k]%=MOD
        return ret

def comb(n,r):
    if n<r:
        return 0
    if n<0 or r<0:
        return 0
    return (fact[n]*pow(fact[r]*fact[n-r],-1,MOD))%MOD

MOD=998244353
N,D=map(int,input().split())
fact=[1]
for i in range(D+1):
    fact.append((fact[-1]*(i+1))%MOD)
ans=0
for i in range(D+2):
    mat=power([[comb(D-1,i),comb(D-1,i-1)],[comb(D-1,i-1),comb(D-1,i-2)]],N)
    ans+=mat[0][0]+mat[1][1]
    ans%=MOD
print(ans)

Submission #46222626 - Tokio Marine & Nichido Fire Insurance Programming Contest 2022（AtCoder Beginner Contest 256)

感想

行列累乗が有効に使えるやつ久しぶりに見た

---

1. 制約より↩