結果

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

ソースコード

diff #

import sys

def main():
    n, K = map(int, sys.stdin.readline().split())
    A = list(map(int, sys.stdin.readline().split()))
    C = list(map(int, sys.stdin.readline().split()))
    
    # Duplicate the arrays to handle circular intervals
    A += A
    C += C
    
    max_color = 50
    mask = [0] * (max_color + 1)  # 1-based indexing
    for c in range(1, max_color + 1):
        start = max(1, c - K)
        end = min(max_color, c + K)
        for d in range(start, end + 1):
            mask[c] |= 1 << (d - 1)
    
    # Precompute prefix sums
    prefix = [0] * (2 * n + 1)
    for i in range(2 * n):
        prefix[i + 1] = prefix[i] + A[i]
    
    # Initialize DP table
    dp = [[0] * (2 * n) for _ in range(2 * n)]
    for i in range(2 * n):
        c = C[i]
        dp[i][i] = 1 << (c - 1)
    
    max_sum = 0
    
    # Check all single stones first
    for i in range(2 * n):
        current_sum = A[i]
        if current_sum > max_sum:
            max_sum = current_sum
    
    for l in range(2, n + 1):
        for i in range(2 * n - l + 1):
            j = i + l - 1
            current_mask = 0
            for k in range(i, j):
                left = dp[i][k]
                right = dp[k + 1][j]
                merged = 0
                # Check left colors contributing to merged
                c = 1
                while c <= max_color:
                    if (left & (1 << (c - 1))) == 0:
                        c += 1
                        continue
                    if (right & mask[c]) != 0:
                        merged |= 1 << (c - 1)
                    c += 1
                # Check right colors contributing to merged
                c = 1
                while c <= max_color:
                    if (right & (1 << (c - 1))) == 0:
                        c += 1
                        continue
                    if (left & mask[c]) != 0:
                        merged |= 1 << (c - 1)
                    c += 1
                current_mask |= merged
            dp[i][j] = current_mask
            if current_mask != 0:
                current_sum = prefix[j + 1] - prefix[i]
                if current_sum > max_sum:
                    max_sum = current_sum
    
    # Check all intervals of length <=n in the duplicated array
    for i in range(2 * n):
        for j in range(i, min(i + n, 2 * n)):
            if dp[i][j] != 0:
                current_sum = prefix[j + 1] - prefix[i]
                if current_sum > max_sum:
                    max_sum = current_sum
    
    print(max_sum)

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