結果

問題 No.3119 A Little Cheat
ユーザー Mistletoe
提出日時 2025-04-19 19:54:21
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
AC  
実行時間 1,068 ms / 2,000 ms
コード長 2,490 bytes
コンパイル時間 407 ms
コンパイル使用メモリ 12,416 KB
実行使用メモリ 162,480 KB
最終ジャッジ日時 2025-04-19 19:55:02
合計ジャッジ時間 35,974 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 49
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys

def solve():
    data = sys.stdin.read().split()
    it = iter(data)
    N = int(next(it)); M = int(next(it))
    A = [int(next(it)) for _ in range(N)]
    mod = 998244353

    # Precompute M^N and M^(N-1)
    powM_N = pow(M, N, mod)
    powM_N1 = pow(M, N-1, mod)

    # Sum of f_no_swap over all B
    sum_no_swap = sum((M - ai) % mod for ai in A) * powM_N1 % mod

    # Helpers to get Xspan and Zspan intervals for each adjacent pair
    def get_spans(ai, aj):
        X, Z = [], []
        if ai < aj:
            # forbidden when B_i ∉ (ai,aj] and B_{i+1} ∈ (ai,aj]
            if ai >= 1:      X.append((1, ai))
            if aj < M:      X.append((aj+1, M))
            if ai+1 <= aj: Z.append((ai+1, aj))
        elif ai > aj:
            # forbidden when B_i ∈ (aj,ai] and B_{i+1} ∉ (aj,ai]
            X.append((aj+1, ai))
            if aj >= 1:      Z.append((1, aj))
            if ai < M:      Z.append((ai+1, M))
        return X, Z

    def span_size(spans):
        total = 0
        for l, r in spans:
            if l <= r:
                total += r-l+1
        return total

    def span_overlap(s1, s2):
        ov = 0
        for l1, r1 in s1:
            for l2, r2 in s2:
                lo = max(l1, l2); hi = min(r1, r2)
                if lo <= hi:
                    ov += hi - lo + 1
        return ov

    # Precompute sizes and overlaps
    size_X = [0]*N
    size_Z = [0]*N
    Xsp = [None]*N
    Zsp = [None]*N
    for i in range(N-1):
        X, Z = get_spans(A[i], A[i+1])
        Xsp[i], Zsp[i] = X, Z
        size_X[i] = span_size(X)
        size_Z[i] = span_size(Z)

    overlap = [0]*N
    for i in range(1, N-1):
        overlap[i] = span_overlap(Xsp[i], Zsp[i-1])

    # S[i] = # of valid prefixes of length i
    # C[i] = sum of DP_i[x] over x in Xspan at position i
    S = [0]*(N+1)
    C = [0]*(N+1)

    # Base: length-1 prefixes
    S[1] = M % mod
    # Compute C[1] and S[2]
    C[1] = size_X[0] % mod
    if N >= 2:
        S[2] = (M * S[1] - C[1] * size_Z[0]) % mod

    # Iterate to build up to length N
    for i in range(2, N):
        C[i]   = (S[i-1] * size_X[i-1] - C[i-1] * overlap[i-1]) % mod
        S[i+1] = (M * S[i]         - C[i] * size_Z[i-1]) % mod

    # T0 = number of B with no beneficial swap
    T0 = S[N] % mod

    # Sequences with any beneficial swap
    total = powM_N
    S1 = (total - T0) % mod

    # Final answer
    print((sum_no_swap + S1) % mod)

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