Official

E - 宝石のネックレス / Jewel Necklace Editorial by admin

Gemini 3.0 Flash (Thinking)

概要

\(N\) 個の宝石から \(K\) 個を選び、それらを円環状に並べたときの「輝きの総和」と「隣り合う宝石同士の調和度の総和」の合計を最大化する問題です。 宝石の数 \(N\) が最大 16 と非常に小さいため、ビット DP(Bitmask DP)を用いて効率的に解くことができます。

考察

1. 問題の言い換え

この問題は、グラフ理論における「重み付きハミルトンサイクル」の探索に似ています。 - 各宝石を頂点とみなし、宝石 \(i\) には重み \(P_i\) があります。 - 宝石 \(i\)\(j\) の間の辺には重み \(C_{i,j}\) があります。 - \(K\) 個の頂点を選び、それらを結ぶ長さ \(K\) のサイクルの(頂点の重み + 辺の重み)を最大化します。

2. なぜ素直な探索ではダメなのか

\(N\) 個から \(K\) 個を選ぶ組み合わせは \(\binom{N}{K}\) 通りあり、さらにそれらを並べる順序(円順列)は \((K-1)! / 2\) 通りあります。 \(N=16, K=16\) の場合、組み合わせは 1 通りですが、並べ方は \(15! / 2 \approx 6.5 \times 10^{11}\) 通りとなり、全探索では制限時間に間に合いません。

3. 動的計画法(ビット DP)の活用

「どの宝石をすでに選んだか」という情報をビット列(mask)で管理することで、計算量を劇的に削減できます。 円環状に並べるため、「どこか 1 つの宝石を始点として固定し、そこからパスを伸ばして最後に始点に戻る」と考えます。

重複を避けるため、各 mask において 「含まれている宝石の中で最も番号が小さいもの」を始点 \(s\) として固定します。これにより、同じサイクルを異なる始点から二重に数えることを防げます。

アルゴリズム

1. DP 状態の定義

dp[mask][i]: - mask: 現在までに選んだ宝石の集合(ビット表現) - i: 最後に選んだ宝石の番号 - 値: その状態における(輝きの総和 + パス上の調和度の総和)の最大値

2. 初期状態

各宝石 \(i\) について、その宝石のみを選んだ状態を初期値とします。 - dp[1 << i][i] = P[i]

3. 遷移

現在の集合 mask、最後に選んだ宝石 i、始点 smask 内の最小インデックス)に対し、次に選ぶ宝石 j を検討します。 - 条件: jmask に含まれておらず、かつ j > s であること(始点固定のため)。 - dp[mask | (1 << j)][j] = max(dp[mask | (1 << j)][j], dp[mask][i] + C[i][j] + P[j])

4. サイクルの完成(答えの計算)

選んだ宝石の数(mask の立っているビット数)がちょうど \(K\) 個になったとき、最後に選んだ宝石 i から始点 s へ戻る調和度を加算します。 - beauty = dp[mask][i] + C[i][s] この beauty の最大値が答えとなります。

計算量

  • 時間計算量: \(O(2^N \cdot N^2)\)
    • 状態数が \(O(2^N \cdot N)\) であり、各状態からの遷移に \(O(N)\) かかるためです。
    • \(N=16\) のとき \(2^{16} \cdot 16^2 \approx 6.5 \times 10^4 \cdot 256 \approx 1.6 \times 10^7\) となり、Python でも十分に高速な実装(ビット演算の最適化など)を行えば制限時間内に収まります。
  • 空間計算量: \(O(2^N \cdot N)\)
    • DP テーブルの大きさに依存します。

実装のポイント

  • 始点の固定: mask の中で最も右にある立っているビット(mask & -mask)を始点 s とすることで、計算の重複を効率よく防いでいます。

  • ループの順序: 宝石の数(popcount)が少ない順に DP を更新していく必要があります。

  • 高速化: Python では bin(mask).count('1') やビット演算を駆使し、内側のループを極力シンプルに保つことが重要です。また、masks_by_size のように、ビット数ごとに mask をあらかじめ分類しておくと、無駄なループを減らせます。

    ソースコード

import sys

def solve():
    # Use fast I/O to read all input data at once
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # Read N and K
    N = int(input_data[0])
    K = int(input_data[1])
    
    # Read jewel brightness values P
    P = list(map(int, input_data[2:2+N]))
    
    # Read harmony values C as an N x N matrix
    C = []
    for i in range(N):
        row = list(map(int, input_data[2+N+i*N:2+N+(i+1)*N]))
        C.append(row)

    # Precompute masks grouped by their popcount and identify the smallest set bit index
    masks_by_size = [[] for _ in range(K + 1)]
    trailing_zeros = [0] * (1 << N)
    bit_to_idx = [0] * (1 << N)
    for i in range(N):
        bit_to_idx[1 << i] = i
    
    for mask in range(1, 1 << N):
        # Calculate the number of set bits (popcount)
        cnt = bin(mask).count('1')
        
        # Identify the smallest index (trailing zero) in the mask to fix the cycle's start
        lowbit = mask & -mask
        tz = bit_to_idx[lowbit]
        trailing_zeros[mask] = tz
        
        # Only store masks with popcount up to K
        if cnt <= K:
            masks_by_size[cnt].append(mask)

    # INF is used for initializing the DP table
    INF = 10**18
    # dp[mask][i] stores the maximum beauty of a path visiting jewels in 'mask' and ending at jewel 'i'.
    # We fix the starting jewel as the one with the smallest index in the mask to avoid redundant calculations.
    dp = [[-INF] * N for _ in range(1 << N)]
    
    # Base case: paths consisting of a single jewel
    for i in range(N):
        dp[1 << i][i] = P[i]

    # Fill the DP table using the bitmask approach for paths of length 1 up to K-1
    for size in range(1, K):
        for mask in masks_by_size[size]:
            s = trailing_zeros[mask]
            dp_mask = dp[mask]
            
            # Identify jewels 'i' currently in the mask that have been reached
            in_mask = []
            temp_mask = mask
            while temp_mask:
                i_bit = temp_mask & -temp_mask
                i = bit_to_idx[i_bit]
                if dp_mask[i] > -INF // 2:
                    in_mask.append(i)
                temp_mask ^= i_bit
            
            if not in_mask:
                continue
            
            # Try adding a new jewel 'j' to the selection.
            # To ensure uniqueness, only consider jewels with an index greater than the fixed start 's'.
            possible_j_mask = ((1 << N) - (1 << (s + 1))) & ~mask
            temp_j_mask = possible_j_mask
            while temp_j_mask:
                j_bit = temp_j_mask & -temp_j_mask
                j = bit_to_idx[j_bit]
                temp_j_mask ^= j_bit
                
                # Pre-fetch the DP list for the updated mask to speed up access
                new_dp_mask = dp[mask | (1 << j)]
                
                # Find the maximum beauty by connecting some 'i' in the mask to 'j'
                max_val = -INF
                for i in in_mask:
                    val = dp_mask[i] + C[i][j]
                    if val > max_val:
                        max_val = val
                
                # Update the DP state for the new mask and end jewel 'j'
                new_val = max_val + P[j]
                if new_val > new_dp_mask[j]:
                    new_dp_mask[j] = new_val

    # Determine the maximum beauty by completing the Hamiltonian cycle for each mask of size K
    ans = -INF
    for mask in masks_by_size[K]:
        s = trailing_zeros[mask]
        dp_mask = dp[mask]
        # Any jewel 'i' in the mask (other than 's') can be the last one before closing the circle back to 's'
        temp_mask = mask ^ (1 << s)
        while temp_mask:
            i_bit = temp_mask & -temp_mask
            i = bit_to_idx[i_bit]
            temp_mask ^= i_bit
            
            if dp_mask[i] > -INF // 2:
                # Add the harmony between the final jewel 'i' and the initial jewel 's'
                val = dp_mask[i] + C[i][s]
                if val > ans:
                    ans = val
    
    # Output the maximum beauty found
    print(ans)

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-thinking によって生成されました。

posted:
last update: