結果

問題 No.840 ほむほむほむら
コンテスト
ユーザー maspy
提出日時 2020-01-02 16:48:46
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
AC  
実行時間 992 ms / 4,000 ms
コード長 1,438 bytes
コンパイル時間 197 ms
コンパイル使用メモリ 12,672 KB
実行使用メモリ 44,668 KB
最終ジャッジ日時 2024-11-22 18:12:29
合計ジャッジ時間 18,982 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 25
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

import itertools
import numpy as np

"""
・empty, ほ、ほむ、ほむら　の個数 mod K の分布ごとに、文字列数を持つ。size <= 125の行列
・遷移は行列で書ける
"""

N,K = map(int,read().split())

MOD = 998244353

def get_index(ho,homu,homura):
    return (ho + homu * K) * K + homura

def create_mat(K):
    L = K ** 3
    A = np.zeros((L,L),np.int64)
    for a,b,c in itertools.product(range(K), repeat = 3):
        i = get_index(a,b,c)
        # ほ
        j = get_index((a+1)%K,b,c)
        A[j,i] += 1
        # む
        j = get_index(a,(a+b)%K,c)
        A[j,i] += 1
        # ら
        j = get_index(a,b,(b+c)%K)
        A[j,i] += 1
    return A

def mult(A,B,MOD):
    mask = (1 << 15) - 1
    A1 = A >> 15; A2 = A & mask
    B1 = B >> 15; B2 = B & mask
    X = np.dot(A1,B1) % MOD
    Y = np.dot(A2,B2) % MOD
    Z = (np.dot(A1+A2,B1+B2) - X - Y) % MOD
    C = (X << 30) + (Z << 15) + Y
    return C % MOD

def power_mat(A,n,MOD):
    k = A.shape[0]
    if n == 0:
        return np.eye(k,dtype=np.int64)
    B = power_mat(A,n//2,MOD)
    B = mult(B,B,MOD)
    return mult(A,B,MOD) if n & 1 else B

A = create_mat(K)
B = power_mat(A,N,MOD)

answer = 0
for a,b in itertools.product(range(K),repeat=2):
    i = get_index(a,b,0)
    answer += B[i,0]
answer %= MOD

print(answer)
0