結果

問題 No.2423 Merge Stones
ユーザー lam6er
提出日時 2025-04-15 22:05:39
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,529 bytes
コンパイル時間 166 ms
コンパイル使用メモリ 81,852 KB
実行使用メモリ 68,776 KB
最終ジャッジ日時 2025-04-15 22:07:22
合計ジャッジ時間 6,826 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 10 TLE * 1 -- * 61
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    idx = 0
    N = int(data[idx])
    idx += 1
    K = int(data[idx])
    idx += 1
    
    A = list(map(int, data[idx:idx+N]))
    idx += N
    C = list(map(int, data[idx:idx+N]))
    idx += N
    
    # Double the arrays to handle circular intervals
    A_doubled = A * 2
    C_doubled = C * 2
    
    # Precompute compatible colors for each color
    compatible = [[] for _ in range(51)]
    for c in range(1, 51):
        min_c = max(1, c - K)
        max_c = min(50, c + K)
        compatible[c] = list(range(min_c, max_c + 1))
    
    # Initialize DP
    INF = -1 << 60
    size = 2 * N
    dp = [[[INF] * 51 for _ in range(size)] for _ in range(size)]
    
    for i in range(size):
        c = C_doubled[i]
        dp[i][i][c] = A_doubled[i]
    
    # Precompute prefix sums
    prefix = [0] * (size + 1)
    for i in range(size):
        prefix[i+1] = prefix[i] + A_doubled[i]
    
    # Fill DP for intervals of length >=2
    for l in range(2, N+1):
        for i in range(size - l + 1):
            j = i + l - 1
            # Iterate all possible splits
            for k in range(i, j):
                # Left is [i..k], right is [k+1..j]
                # Check all compatible color pairs
                # Iterate over left's colors and compatible colors in right
                for c1 in range(1, 51):
                    if dp[i][k][c1] <= INF:
                        continue
                    # Iterate over compatible colors for c1
                    for c2 in compatible[c1]:
                        if c2 < 1 or c2 > 50:
                            continue
                        if dp[k+1][j][c2] <= INF:
                            continue
                        total = dp[i][k][c1] + dp[k+1][j][c2]
                        # Update merged interval's c1 and c2
                        if total > dp[i][j][c1]:
                            dp[i][j][c1] = total
                        if total > dp[i][j][c2]:
                            dp[i][j][c2] = total
    
    # Find the maximum sum across all intervals of length <= N
    max_sum = 0
    for i in range(size):
        for j in range(i, min(i + N, size)):
            current_length = j - i + 1
            if current_length > N:
                continue
            for c in range(1, 51):
                if dp[i][j][c] > max_sum:
                    max_sum = dp[i][j][c]
    
    print(max_sum)

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