結果

問題 No.2423 Merge Stones
コンテスト
ユーザー gew1fw
提出日時 2025-06-12 18:39:34
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,298 bytes
コンパイル時間 180 ms
コンパイル使用メモリ 81,924 KB
実行使用メモリ 67,992 KB
最終ジャッジ日時 2025-06-12 18:39:41
合計ジャッジ時間 6,628 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 10 TLE * 1 -- * 61
権限があれば一括ダウンロードができます

ソースコード

diff # 

def main():
    import sys
    input = sys.stdin.read().split()
    idx = 0
    N = int(input[idx])
    idx += 1
    K = int(input[idx])
    idx += 1
    A = list(map(int, input[idx:idx+N]))
    idx += N
    C = list(map(int, input[idx:idx+N]))
    idx += N

    # Double the arrays to handle circularity
    A = A + A
    C = C + C
    max_len = N
    total_size = 2 * N

    # Precompute prefix sums
    prefix = [0] * (total_size + 1)
    for i in range(total_size):
        prefix[i+1] = prefix[i] + A[i]

    # Precompute S[i][j]: set of colors in C[i..j]
    S = [[set() for _ in range(total_size)] for __ in range(total_size)]
    for i in range(total_size):
        current = set()
        for j in range(i, min(i + max_len, total_size)):
            current.add(C[j])
            S[i][j] = current.copy()

    # Precompute compatible colors
    max_color = max(C)
    compatible = [[False] * (max_color + 1) for _ in range(max_color + 1)]
    for c1 in range(max_color + 1):
        for c2 in range(max_color + 1):
            if abs(c1 - c2) <= K:
                compatible[c1][c2] = True

    # DP table: dp[i][j] is a set of possible colors
    dp = [[set() for _ in range(total_size)] for __ in range(total_size)]

    # Initialize base case
    for i in range(total_size):
        dp[i][i].add(C[i])

    # Fill DP table
    for l in range(2, max_len + 1):
        for i in range(total_size - l + 1):
            j = i + l - 1
            if j >= total_size:
                continue
            current = set()
            for k in range(i, j):
                left = dp[i][k]
                right = dp[k+1][j]
                for c1 in left:
                    for c2 in right:
                        if compatible[c1][c2]:
                            current.add(c1)
                            current.add(c2)
            # Restrict to colors in S[i][j]
            current = {c for c in current if c in S[i][j]}
            dp[i][j] = current

    # Find the maximum sum
    max_sum = 0
    for i in range(total_size):
        for j in range(i, min(i + max_len, total_size)):
            if len(dp[i][j]) > 0:
                s = prefix[j+1] - prefix[i]
                if s > max_sum:
                    max_sum = s
    print(max_sum)

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