結果

問題 No.2423 Merge Stones
ユーザー qwewe
提出日時 2025-05-14 13:22:59
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 3,147 bytes
コンパイル時間 638 ms
コンパイル使用メモリ 82,912 KB
実行使用メモリ 128,484 KB
最終ジャッジ日時 2025-05-14 13:25:04
合計ジャッジ時間 6,642 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 10 TLE * 1 -- * 61
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys

def solve():
    N, K = map(int, sys.stdin.readline().split())
    A = list(map(int, sys.stdin.readline().split()))
    C = list(map(int, sys.stdin.readline().split()))

    # Using -1 to represent "not possible" since A_i >= 1, so valid sums are positive.
    neg_inf = -1 

    # dp[length][start_idx][color_idx]
    # color_idx is 0-indexed (color_value - 1)
    # MaxColor is 50. Indices 0 to 49.
    dp = [[[neg_inf] * 50 for _ in range(N)] for _ in range(N + 1)]

    max_overall_power = 0 # Since all A_i >= 1, min possible positive power is 1.

    # Base cases: segments of length 1
    for i in range(N):
        color_val = C[i]
        color_idx = color_val - 1 # Convert to 0-indexed
        dp[1][i][color_idx] = A[i]
        if A[i] > max_overall_power:
            max_overall_power = A[i]

    # Iterate over segment lengths
    for length in range(2, N + 1):
        # Iterate over starting index of the segment
        for i in range(N): 
            # Iterate over split point (k_len is length of the left part)
            for k_len in range(1, length): 
                left_start_idx = i
                left_len = k_len
                
                right_start_idx = (i + k_len) % N
                right_len = length - k_len

                # Iterate over color of the left merged part
                for lc_idx in range(50): 
                    val_left = dp[left_len][left_start_idx][lc_idx]
                    if val_left == neg_inf:
                        continue
                    
                    # Iterate over color of the right merged part,
                    # ensuring it's compatible with left_color.
                    # abs(left_color_val - right_color_val) <= K
                    # abs((lc_idx+1) - (rc_idx+1)) <= K  => abs(lc_idx - rc_idx) <= K
                    # So, rc_idx must be in [lc_idx - K, lc_idx + K]
                    
                    min_compatible_rc_idx = max(0, lc_idx - K)
                    max_compatible_rc_idx = min(49, lc_idx + K)

                    for rc_idx in range(min_compatible_rc_idx, max_compatible_rc_idx + 1):
                        val_right = dp[right_len][right_start_idx][rc_idx]
                        if val_right == neg_inf:
                            continue
                        
                        current_sum = val_left + val_right
                        
                        # The new merged stone (segment of `length` starting at `i`)
                        # can take color of left part (lc_idx) or right part (rc_idx).
                        
                        if current_sum > dp[length][i][lc_idx]:
                            dp[length][i][lc_idx] = current_sum
                        
                        if current_sum > dp[length][i][rc_idx]:
                            dp[length][i][rc_idx] = current_sum
                        
                        # Update overall maximum power found
                        if current_sum > max_overall_power:
                            max_overall_power = current_sum
                            
    print(max_overall_power)

solve()
0