概要

重さの制限 \(W\) を満たしつつ、「選んだ展示品の利益の合計」から「選んだ会場の種類数 \(c\) の2乗」を引いた値を最大化する、ナップサック問題の応用問題です。

考察

この問題は、通常のナップサック問題に「選んだ展示品の会場の種類数に応じたペナルティ（配送コスト）」が追加された形をしています。

まず、制約の \(K \le 8\) （会場の数が非常に少ない） という点に注目します。 展示品を1つずつ順番に見て「選ぶ・選ばない」を決める標準的なDPを行う場合、「どの会場の展示品を選んだか」を記憶しておく必要があります。これを集合（ビットマスク）で持つと、状態数は \(2^K \times W\) となり、計算量は \(O(N \cdot 2^K \cdot W)\) となります。これでも解けなくはないですが、Pythonなどの言語では実行時間制限（TLE）に引っかかる恐れがあります。

そこで、展示品を「会場ごと」にグループ化して処理する ことを考えます。 会場ごとにまとめて「この会場からは1つも選ばない」か「この会場から1つ以上選ぶ」かを決めていけば、具体的な会場の組み合わせを覚えておく必要はなく、「これまでにいくつの会場から選んだか（種類数 \(c\)）」だけを状態として持てばよくなります。これにより、状態数を \(K \times W\) に減らすことができ、より高速に解くことができます。

アルゴリズム

1. グループ化 すべての展示品を、出展先の会場ごとにグループ分けします。

2. DPテーブルの定義 \(dp[c][w]\) を、「合計 \(c\) 種類の会場から展示品を選び、重さの合計が \(w\) であるときの最大利益」と定義します。初期状態は \(dp[0][0] = 0\) で、それ以外は非常に小さい値（\(-\infty\)）にしておきます。

3. DPの遷移 各会場ごとに順番にDPテーブルを更新していきます。ある会場の展示品たちを見るとき、以下の2つのケースを考えます。

   • その会場から1つも選ばない場合： 前の会場までの状態がそのまま引き継がれます。
   • その会場から1つ以上選ぶ場合： 選んだ会場の種類数 \(c\) が \(1\) 増えます。そのため、前の会場までの状態 \(dp[c-1]\) をベースにして、現在の会場の展示品だけでナップサックDPを行います。

4. 答えの計算 すべての会場を見終わった後の \(dp[c][w]\) には、各条件での最大利益が入っています。 すべての \(c\) (\(0 \le c \le K\)) と \(w\) (\(0 \le w \le W\)) の組み合わせについて、\(dp[c][w] - c^2\) を計算し、その中の最大値が答えとなります。

計算量

• 時間計算量: \(O(N \cdot K \cdot W)\) 展示品 \(N\) 個それぞれについて、会場の種類数 \(c\)（最大 \(K\)）と重さ \(w\)（最大 \(W\)）のループを回すため、全体の計算量は \(O(N K W)\) となります。
• 空間計算量: \(O(K \cdot W)\) DPテーブルとして、会場の種類数 \(K\) と重さ \(W\) の2次元配列を保持するため、\(O(K W)\) となります。

実装のポイント

同じ会場から「複数の展示品を選ぶ」場合の処理に少し工夫が必要です。 コード内では、現在の会場から「少なくとも1つ選んだ状態」を管理するための配列 temp_c を用意しています。

現在の会場にある展示品（重さ \(w_i\)、利益 \(v_i\)）を選ぶとき、遷移元（選ぶ前の状態）は以下の2通りあります。 1. この会場から初めて展示品を選ぶ場合： 前の会場までの状態である dp[c-1][w - w_i] に \(v_i\) を足す。 2. すでにこの会場から他の展示品を選んでいる場合： 現在の会場での途中経過である temp_c[w - w_i] に \(v_i\) を足す。

これら2つのうち大きい方を採用して temp_c[w] を更新していくことで、「1つの会場から1つ以上選ぶ」という条件を正しく満たしながら、通常のナップサック問題と同じように処理することができます。また、同じ展示品を複数回選ばないように、重さ \(w\) のループは大きい方から逆順（\(W\) から \(w_i\) まで）に回す必要があります。

ソースコード

import sys

def solve():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    
    N = int(data[0])
    K = int(data[1])
    W = int(data[2])
    
    if N == 0:
        print(0)
        return

    items = []
    idx = 3
    for _ in range(N):
        w = int(data[idx])
        v = int(data[idx+1])
        d = int(data[idx+2])
        items.append((w, v, d))
        idx += 3
        
    venues = list(set(d for w, v, d in items))
    M = len(venues)
    venue_to_id = {v: i for i, v in enumerate(venues)}
    
    items_by_venue = [[] for _ in range(M)]
    for w, v, d in items:
        items_by_venue[venue_to_id[d]].append((w, v))
        
    INF = 10**15
    # dp[c][w]: c個の異なる会場の展示品を選び、重さの合計がwのときの最大利益
    dp = [[-INF] * (W + 1) for _ in range(M + 1)]
    dp[0][0] = 0
    
    for k in range(M):
        next_dp = [row[:] for row in dp]
        
        for c in range(1, k + 2):
            temp_c = [-INF] * (W + 1)
            dp_prev = dp[c - 1]
            
            for w_i, v_i in items_by_venue[k]:
                for w in range(W, w_i - 1, -1):
                    prev_w = w - w_i
                    val1 = temp_c[prev_w]
                    val2 = dp_prev[prev_w]
                    
                    m = val1 if val1 > val2 else val2
                    if m != -INF:
                        m += v_i
                        if m > temp_c[w]:
                            temp_c[w] = m
                            
            next_dp_c = next_dp[c]
            for w in range(W + 1):
                if temp_c[w] > next_dp_c[w]:
                    next_dp_c[w] = temp_c[w]
                    
        dp = next_dp

    ans = 0
    for c in range(M + 1):
        cost = c * c
        max_v = max(dp[c])
        if max_v - cost > ans:
            ans = max_v - cost
            
    print(ans)

if __name__ == '__main__':
    solve()

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