結果

問題 No.840 ほむほむほむら
ユーザー maspymaspy
提出日時 2020-01-02 16:48:46
言語 Python3
(3.12.2 + numpy 1.26.4 + scipy 1.12.0)
結果
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
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 525 ms
44,028 KB
testcase_01 AC 529 ms
44,020 KB
testcase_02 AC 541 ms
44,124 KB
testcase_03 AC 602 ms
44,420 KB
testcase_04 AC 536 ms
44,288 KB
testcase_05 AC 537 ms
44,024 KB
testcase_06 AC 540 ms
44,280 KB
testcase_07 AC 553 ms
44,656 KB
testcase_08 AC 655 ms
44,532 KB
testcase_09 AC 544 ms
44,280 KB
testcase_10 AC 533 ms
43,904 KB
testcase_11 AC 548 ms
44,540 KB
testcase_12 AC 565 ms
44,404 KB
testcase_13 AC 845 ms
44,280 KB
testcase_14 AC 591 ms
44,408 KB
testcase_15 AC 533 ms
44,412 KB
testcase_16 AC 539 ms
44,536 KB
testcase_17 AC 606 ms
44,032 KB
testcase_18 AC 925 ms
44,024 KB
testcase_19 AC 992 ms
44,440 KB
testcase_20 AC 527 ms
44,276 KB
testcase_21 AC 526 ms
44,280 KB
testcase_22 AC 537 ms
44,028 KB
testcase_23 AC 976 ms
43,900 KB
testcase_24 AC 534 ms
44,416 KB
testcase_25 AC 543 ms
44,668 KB
testcase_26 AC 545 ms
44,404 KB
testcase_27 AC 974 ms
44,536 KB
権限があれば一括ダウンロードができます

ソースコード

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