結果

問題 No.1516 simple 門松列 problem Re:MASTER
コンテスト
ユーザー ansain
提出日時 2021-05-04 03:20:30
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
AC  
実行時間 1,207 ms / 6,000 ms
コード長 2,680 bytes
コンパイル時間 214 ms
コンパイル使用メモリ 12,800 KB
実行使用メモリ 43,820 KB
最終ジャッジ日時 2024-10-09 04:27:58
合計ジャッジ時間 14,677 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 19
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from collections import defaultdict, Counter, deque
from itertools import permutations, combinations, product, combinations_with_replacement, groupby, accumulate
import operator
from math import sqrt, gcd, factorial
# from math import isqrt, prod,comb  # python3.8用(notpypy)
#from bisect import bisect_left,bisect_right
#from functools import lru_cache,reduce
#from heapq import heappush,heappop,heapify,heappushpop,heapreplace
import numpy as np
#import networkx as nx
#from networkx.utils import UnionFind
#from numba import njit, b1, i1, i4, i8, f8
#from scipy.sparse import csr_matrix
#from scipy.sparse.csgraph import shortest_path, floyd_warshall, dijkstra, bellman_ford, johnson, NegativeCycleError
# numba例 @njit(i1(i4[:], i8[:, :]),cache=True) 引数i4配列、i8 2次元配列,戻り値i1
def input(): return sys.stdin.readline().rstrip()
def divceil(n, k): return 1+(n-1)//k  # n/kの切り上げを返す
def yn(hantei, yes='Yes', no='No'): print(yes if hantei else no)


def check(a):
    return a[0] != a[1] and a[1] != a[2] and a[2] != a[0] and a[1] in [max(a), min(a)]


def dot2(mat1, mat2, MOD):
    """
    行列積a@bの10**9+7modをオーバーフロー回避するために上下15bitで分割して掛け算する 
    https://ikatakos.com/pot/programming_algorithm/python_tips/avoid_overflow
    """
    mask = (1 << 15) - 1
    mat1h = mat1 >> 15
    mat1l = mat1 & mask
    mat2h = mat2 >> 15
    mat2l = mat2 & mask
    mathh = mat1h @ mat2h % MOD
    matll = mat1l @ mat2l % MOD
    mathl = (mathh + matll - (mat1h - mat1l) @ (mat2h - mat2l)) % MOD
    res = (mathh << 30) + (mathl << 15) + matll
    res %= MOD
    return res


def matrix_multiplication(a, n, k, mod=10**9+7):  # n次正方行列a^k
    ans = np.eye(n, dtype=np.int64)
    while k:
        k, i = divmod(k, 2)
        if i:
            ans = dot2(ans, a, mod)
        a = dot2(a, a, mod)
    return ans


def main():
    mod = 10**9+7
    mod2 = 998244353
    n, k = map(int, input().split())
    mat = np.zeros((2*k*k, 2*k*k), dtype=np.int64)
    ans = np.zeros(2*k*k, dtype=np.int64)
    for a in range(k):
        for b in range(k):
            for c in range(k):
                if check([a, b, c]):
                    mat[a*k+b][b*k+c] += 1
                    mat[a*k+b+k**2][b*k+c+k**2] += 1
                    mat[a*k+b][b*k+c+k**2] += a
            if a != b:
                ans[a*k+b] += 1
                ans[a*k+b+k**2] += a+b
    mat2 = matrix_multiplication(mat, 2*k*k, n-2, mod2)
    ans2=dot2(ans,mat2,mod2)
    #print(ans2[:k*k])
    #print(ans2[k*k:])
    print(ans2[:k*k].sum()%mod2,ans2[k*k:].sum()%mod2)


if __name__ == '__main__':
    main()
0