結果

問題 No.2423 Merge Stones
ユーザー lam6er
提出日時 2025-04-15 21:59:14
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,686 bytes
コンパイル時間 496 ms
コンパイル使用メモリ 81,816 KB
実行使用メモリ 86,620 KB
最終ジャッジ日時 2025-04-15 22:00:37
合計ジャッジ時間 6,611 ms
ジャッジサーバーID
(参考情報)
judge3 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
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_double = A + A
    C_double = C + C
    max_color = 50

    # Precompute compatible masks for each color (0-based in the array, but colors are 1-based)
    compatible_masks = [0] * max_color  # compatible_masks[c] is for color c+1
    for c in range(max_color):
        current_color = c + 1
        low = max(1, current_color - K)
        high = min(max_color, current_color + K)
        mask = 0
        for d in range(low, high + 1):
            mask |= 1 << (d - 1)
        compatible_masks[c] = mask

    # Precompute prefix sums for the double array
    prefix = [0] * (2 * N + 1)
    for i in range(2 * N):
        prefix[i + 1] = prefix[i] + A_double[i]

    # Initialize possible[i][j] as a list of bitmasks
    possible = [[0] * (2 * N) for _ in range(2 * N)]
    for i in range(2 * N):
        color = C_double[i] - 1  # 0-based
        possible[i][i] = 1 << color

    # Fill the DP table
    for l in range(2, N + 1):
        for i in range(2 * N - l + 1):
            j = i + l - 1
            possible[i][j] = 0
            for k in range(i, j):
                left = possible[i][k]
                right = possible[k + 1][j]
                if left == 0 or right == 0:
                    continue

                # Compute expand_R
                expand_R = 0
                temp_right = right
                c = 0
                while temp_right:
                    if temp_right & 1:
                        expand_R |= compatible_masks[c]
                    temp_right >>= 1
                    c += 1

                # Compute expand_L
                expand_L = 0
                temp_left = left
                c = 0
                while temp_left:
                    if temp_left & 1:
                        expand_L |= compatible_masks[c]
                    temp_left >>= 1
                    c += 1

                merged_mask = (left & expand_R) | (right & expand_L)
                possible[i][j] |= merged_mask

    # Find the maximum sum
    max_sum = 0
    for i in range(2 * N):
        for j in range(i, min(i + N, 2 * N)):
            if j - i + 1 > N:
                continue
            if possible[i][j] == 0:
                continue
            current_sum = prefix[j + 1] - prefix[i]
            if current_sum > max_sum:
                max_sum = current_sum

    print(max_sum)

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