結果

問題 No.2423 Merge Stones
ユーザー lam6er
提出日時 2025-03-31 17:51:11
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,576 bytes
コンパイル時間 265 ms
コンパイル使用メモリ 82,760 KB
実行使用メモリ 67,120 KB
最終ジャッジ日時 2025-03-31 17:52:14
合計ジャッジ時間 7,014 ms
ジャッジサーバーID
（参考情報） 		judge1 / 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
    
    if N == 1:
        print(A[0])
        return
    
    duplicated_A = A * 2
    duplicated_C = C * 2
    max_color = 50
    
    # Precompute prefix sums
    prefix_sum = [0] * (2 * N + 1)
    for i in range(2 * N):
        prefix_sum[i+1] = prefix_sum[i] + duplicated_A[i]
    
    # Precompute allowed masks for each color
    allowed_masks = [0] * (max_color + 1)  # 1-based
    for c in range(1, max_color + 1):
        low = max(1, c - K)
        high = min(max_color, c + K)
        mask = 0
        for neighbor in range(low, high + 1):
            mask |= 1 << (neighbor - 1)  # neighbor is 1-based, stored as bit (neighbor-1)
        allowed_masks[c] = mask
    
    # Initialize DP table
    DP = [[0] * (2 * N) for _ in range(2 * N)]
    for i in range(2 * N):
        color = duplicated_C[i]
        DP[i][i] = 1 << (color - 1)  # color 1 is bit 0
    
    max_sum = max(A)
    
    for length in range(2, N + 1):
        for i in range(2 * N - length + 1):
            j = i + length - 1
            current_mask = 0
            for k in range(i, j):
                left_mask = DP[i][k]
                right_mask = DP[k+1][j]
                
                left_compatible = 0
                for c in range(1, max_color + 1):
                    if (left_mask & (1 << (c-1))) == 0:
                        continue
                    if (allowed_masks[c] & right_mask) != 0:
                        left_compatible |= (1 << (c-1))
                
                right_compatible = 0
                for c in range(1, max_color + 1):
                    if (right_mask & (1 << (c-1))) == 0:
                        continue
                    if (allowed_masks[c] & left_mask) != 0:
                        right_compatible |= (1 << (c-1))
                
                merged = left_compatible | right_compatible
                current_mask |= merged
            
            DP[i][j] = current_mask
            if current_mask != 0:
                current_sum = prefix_sum[j + 1] - prefix_sum[i]
                if current_sum > max_sum:
                    max_sum = current_sum
    
    # Check for all possible intervals in the original circular arrangement by considering all starts
    # Also check all lengths up to N
    print(max_sum)

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