結果
問題 |
No.3119 A Little Cheat
|
ユーザー |
|
提出日時 | 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 |
ソースコード
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()