#1084094 (C++17) No.2423 Merge Stones

• 問題: No.2423 Merge Stones
• ユーザー: qwewe
• 提出日時: 2025-05-14 13:18:47
• 言語: C++17 (gcc 13.3.0 + boost 1.87.0)
• 結果: AC
• 実行時間: 291 ms / 4,000 ms
• コード長: 7,674 bytes
• コンパイル時間: 797 ms
• コンパイル使用メモリ: 84,392 KB
• 実行使用メモリ: 7,844 KB
• 最終ジャッジ日時: 2025-05-14 13:19:47
• 合計ジャッジ時間: 8,661 ms
• ジャッジサーバーID （参考情報）: judge3 / judge4

ファイルパターン	結果
sample	AC * 1
other	AC * 72

ソースコード

#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>

using namespace std;

// Use long long for potentially large sums of magic powers
typedef long long ll;
// Use unsigned long long for bitmasks to represent sets of colors.
// Colors are up to 50, so 64 bits are sufficient.
typedef unsigned long long ull;

// Function to compute the dilated bitmask.
// For each color 'c' represented by a set bit in 'mask', this function sets bits
// for all colors 'c'' such that |c - c'| <= K.
// The resulting mask only includes bits corresponding to valid colors in the range [1, 50].
ull dilate(ull mask, int K) {
    // If K=0, dilation has no effect, return the original mask.
    if (K == 0) return mask;

    // Start with the original mask.
    ull dilated_mask = mask;
    
    // To avoid using already shifted bits for further shifts in the same dilation step,
    // we apply shifts based on the original mask 'mask'.
    // Iterate through shift amounts from 1 to K.
    for (int s = 1; s <= K; ++s) {
        // Shift left by s: colors c become c+s.
        dilated_mask |= (mask << s);
        // Shift right by s: colors c become c-s.
        dilated_mask |= (mask >> s);
    }
    
    // Create a mask to keep only bits corresponding to colors 1 through 50.
    // (1ULL << 51) represents 2^51, which has bit 51 set.
    // (1ULL << 1) represents 2^1, which has bit 1 set.
    // The difference (2^51 - 2^1) creates a mask with bits 1 to 50 set.
    // Example: (1<<4) - (1<<1) = 16 - 2 = 14 = binary 1110 (bits 1, 2, 3 set).
    ull valid_mask = (1ULL << 51) - (1ULL << 1);
    
    // Apply the valid mask to clear any bits outside the range [1, 50].
    return dilated_mask & valid_mask;
}


int main() {
    // Use fast I/O operations.
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int N, K;
    cin >> N >> K;

    vector<ll> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
    }

    vector<int> C(N);
    for (int i = 0; i < N; ++i) {
        cin >> C[i];
    }

    // Double the arrays A and C to handle the circular arrangement seamlessly.
    // A stone originally at index i (0 to N-1) is now at indices i and i+N.
    vector<ll> A_double(2 * N);
    vector<int> C_double(2 * N);
    for (int i = 0; i < N; ++i) {
        A_double[i] = A_double[i + N] = A[i];
        C_double[i] = C_double[i + N] = C[i];
    }

    // Compute prefix sums for the doubled magic power array A_double.
    // P[k] will store the sum of A_double[0]...A_double[k-1].
    // P has size 2*N + 1 to accommodate sums up to index 2*N - 1. P[0] = 0.
    vector<ll> P(2 * N + 1, 0);
    for (int i = 0; i < 2 * N; ++i) {
        P[i + 1] = P[i] + A_double[i];
    }

    // DP table definition. dp[i][j] stores a bitmask representing the set of possible colors
    // that a single merged stone covering the original stones corresponding to indices i..j 
    // (in the doubled array) can have.
    // The table size is 2N x 2N. Indices are 0-based.
    vector<vector<ull>> dp(2 * N, vector<ull>(2 * N, 0));

    // Base cases: Blocks of length 1 correspond to individual stones.
    // The only possible color is the stone's own color.
    for (int i = 0; i < 2 * N; ++i) {
        // Check if the color is within the valid range [1, 50].
        if (C_double[i] >= 1 && C_double[i] <= 50) {
            // Set the bit corresponding to the color C_double[i].
             dp[i][i] = (1ULL << C_double[i]);
        }
    }

    // Dynamic Programming: Fill the table for blocks of length 2 up to N.
    // Iterate through lengths `len` from 2 to N.
    for (int len = 2; len <= N; ++len) {
        // Iterate through all possible starting indices `i` for a block of length `len`.
        // The starting index `i` can range from 0 up to 2*N - len.
        for (int i = 0; i <= 2 * N - len; ++i) {
            int j = i + len - 1; // Ending index of the block [i..j].
            
            // Initialize the dp state for block [i..j] to 0. It will be updated if merges are possible.
            dp[i][j] = 0; 

            // Iterate through all possible split points `k`.
            // A split at `k` divides the block [i..j] into two sub-blocks: [i..k] and [k+1..j].
            for (int k = i; k < j; ++k) {
                ull bm1 = dp[i][k]; // Bitmask of possible colors for the left sub-block [i..k].
                ull bm2 = dp[k + 1][j]; // Bitmask of possible colors for the right sub-block [k+1..j].

                // If either sub-block cannot be merged into a single stone (mask is 0),
                // then this split point `k` cannot lead to merging the block [i..j]. Skip.
                if (bm1 == 0 || bm2 == 0) {
                    continue; 
                }

                // Compute the dilated masks. d_bm1 contains colors within K distance of colors in bm1.
                // d_bm2 contains colors within K distance of colors in bm2.
                ull d_bm1 = dilate(bm1, K);
                ull d_bm2 = dilate(bm2, K);

                // Check if a merge is possible between the two sub-blocks.
                // This requires finding at least one pair of colors (c1, c2) such that
                // c1 is possible for block [i..k] (c1 is in bm1),
                // c2 is possible for block [k+1..j] (c2 is in bm2),
                // and |c1 - c2| <= K.
                // This condition is efficiently checked by intersecting bm1 with the dilated mask of bm2.
                // If the intersection is non-zero, at least one such pair exists.
                if ((bm1 & d_bm2) != 0) {
                    // If a merge is possible, update the dp state for block [i..j].
                    // The resulting merged stone can take the color of either stone involved in the final merge.
                    // If the merge was possible due to pair (c1, c2), the resulting color could be c1 or c2.
                    // The set of all possible c1's that participate in *any* valid merge via split k is (bm1 & d_bm2).
                    // The set of all possible c2's that participate in *any* valid merge via split k is (bm2 & d_bm1).
                    // The final set of possible colors for block [i..j] resulting from *this split* k is the union of these two sets.
                    // We accumulate possibilities from all valid splits using bitwise OR.
                    dp[i][j] |= (bm1 & d_bm2) | (bm2 & d_bm1);
                }
            }
        }
    }

    // After filling the DP table, find the maximum magic power among all possible merged stones.
    ll max_power = 0;
    // Iterate through all possible block lengths from 1 to N.
    for (int len = 1; len <= N; ++len) {
         // Iterate through all possible starting positions `i` corresponding to the original N stones (indices 0 to N-1).
        for (int i = 0; i < N; ++i) { 
            int j = i + len - 1; // Calculate the ending index in the doubled array.
            
            // Check if the block corresponding to indices i..j in the doubled array can be merged into a single stone.
            // This is true if dp[i][j] is non-zero (has at least one possible color).
            if (dp[i][j] != 0) {
                // If mergeable, calculate the total magic power of this block using prefix sums.
                // The sum of A_double[i] through A_double[j] is P[j+1] - P[i].
                ll current_power = P[j + 1] - P[i];
                // Update the overall maximum power found.
                max_power = max(max_power, current_power);
            }
        }
    }
    
    // Output the final maximum possible magic power.
    cout << max_power << endl;

    return 0;
}