結果

問題 No.1513 simple 門松列 problem
コンテスト
ユーザー lam6er
提出日時 2025-04-16 00:22:17
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 713 ms / 3,000 ms
コード長 3,317 bytes
コンパイル時間 365 ms
コンパイル使用メモリ 81,860 KB
実行使用メモリ 101,216 KB
最終ジャッジ日時 2025-04-16 00:23:43
合計ジャッジ時間 5,260 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 18
権限があれば一括ダウンロードができます

ソースコード

diff # 

MOD = 998244353

def main():
    import sys
    input = sys.stdin.read().split()
    N = int(input[0])
    K = int(input[1])
    
    if N < 3:
        print(0, 0)
        return
    
    # Initialize for i=2
    prev_count = [[0] * K for _ in range(K)]
    prev_sum = [[0] * K for _ in range(K)]
    for a in range(K):
        for b in range(K):
            if a != b:
                prev_count[a][b] = 1
                prev_sum[a][b] = a + b
    
    for i in range(3, N + 1):
        current_count = [[0] * K for _ in range(K)]
        current_sum = [[0] * K for _ in range(K)]
        # Precompute prefix sums for each b_prev
        prefix_count = [[0] * K for _ in range(K)]
        prefix_sum = [[0] * K for _ in range(K)]
        for b_prev in range(K):
            # Compute prefix_count[b_prev][x]
            prefix_count[b_prev][0] = prev_count[0][b_prev]
            prefix_sum[b_prev][0] = prev_sum[0][b_prev]
            for x in range(1, K):
                prefix_count[b_prev][x] = (prefix_count[b_prev][x-1] + prev_count[x][b_prev]) % MOD
                prefix_sum[b_prev][x] = (prefix_sum[b_prev][x-1] + prev_sum[x][b_prev]) % MOD
        
        # Iterate over all (b_prev, c)
        for b_prev in range(K):
            for c in range(K):
                if b_prev == c:
                    continue
                sum_count = 0
                sum_sum = 0
                if b_prev > c:
                    # a_prev in 0..b_prev-1, exclude c
                    if b_prev == 0:
                        sum_count = 0
                        sum_sum = 0
                    else:
                        sum_count = prefix_count[b_prev][b_prev - 1]
                        sum_sum = prefix_sum[b_prev][b_prev - 1]
                        # Subtract if c is within the range
                        if c < b_prev:
                            sum_count = (sum_count - prev_count[c][b_prev]) % MOD
                            sum_sum = (sum_sum - prev_sum[c][b_prev]) % MOD
                else:
                    # b_prev < c: a_prev in b_prev+1..K-1, exclude c
                    if b_prev + 1 >= K:
                        sum_count = 0
                        sum_sum = 0
                    else:
                        sum_count = (prefix_count[b_prev][K-1] - prefix_count[b_prev][b_prev]) % MOD
                        sum_sum = (prefix_sum[b_prev][K-1] - prefix_sum[b_prev][b_prev]) % MOD
                        # Subtract if c is in the range
                        if c >= b_prev + 1 and c < K:
                            sum_count = (sum_count - prev_count[c][b_prev]) % MOD
                            sum_sum = (sum_sum - prev_sum[c][b_prev]) % MOD
                # Update current_count and current_sum
                current_count[b_prev][c] = (current_count[b_prev][c] + sum_count) % MOD
                current_sum[b_prev][c] = (current_sum[b_prev][c] + (sum_sum + sum_count * c) % MOD) % MOD
        
        # Update prev_count and prev_sum for next iteration
        prev_count, prev_sum = current_count, current_sum
    
    ans_1 = 0
    ans_2 = 0
    for a in range(K):
        for b in range(K):
            ans_1 = (ans_1 + prev_count[a][b]) % MOD
            ans_2 = (ans_2 + prev_sum[a][b]) % MOD
    print(ans_1, ans_2)

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