概要

\(N\) 人の社員からちょうど \(K\) 人を選んでチームを作るとき、「選んだ社員の能力値の合計」から「チーム内に含まれる相性の悪いペアのペナルティ合計」を引いた値を最大化する問題です。

考察

この問題を解くための鍵は、制約の小ささに注目することです。

• 社員数 \(N\) が最大 18 と非常に小さい
• \(N\) 人から \(K\) 人を選ぶ組み合わせの数 \(\binom{N}{K}\) は、最も多くなる \(N=18, K=9\) の場合でも \(\binom{18}{9} = 48,620\) 通りしかありません。

このように組み合わせの数が十分に少ない場合、すべての組み合わせを力まかせに調べる「全探索」という手法が有効です。

各組み合わせに対して、以下の計算を行います： 1. 選んだ \(K\) 人の能力値 \(A_i\) をすべて足す。 2. 選んだ \(K\) 人の中に相性の悪いペア \((U_j, V_j)\) が含まれていれば、そのペナルティ \(B_j\) を引く。

この計算をすべての組み合わせに対して行い、その中の最大値が答えとなります。

アルゴリズム

1. ペナルティ情報の整理: 社員 \(u\) と社員 \(v\) が同時に選ばれた時の減少量 \(B\) を、二次元配列 penalty_matrix[u][v] に保存しておきます。これにより、特定のペアが相性が悪いかどうかを \(O(1)\) ですぐに確認できるようになります。

2. 組み合わせの列挙: Python の itertools.combinations を使い、\(N\) 人から \(K\) 人を選ぶ全パターンを列挙します。

3. 総合力の計算: 各組み合わせについて、以下の手順でスコアを計算します。

   • チームメンバーを \(u\) としたとき、current_score += A[u] を行う。
   • さらに、チーム内の別のメンバー \(v\) との間にペナルティがあるか確認し、あれば current_score -= penalty_matrix[u][v] を行う。
   • このとき、同じペアを 2 回計算しないように注意します。

4. 最大値の更新: 計算したスコアがこれまでの最大値よりも大きければ、値を更新します。

計算量

• 時間計算量: \(O(N^2 + \binom{N}{K} \cdot K^2)\)

  • ペナルティ行列の初期化に \(O(N^2)\) かかります。
  • 組み合わせの数が \(\binom{N}{K}\) 通りあり、各組み合わせにつき \(K\) 人の能力値加算と、約 \(K^2/2\) 回のペナルティ判定を行うため、全体でこの計算量になります。\(N=18\) のとき、最悪でも数百万回程度の計算で済むため、制限時間内に十分間に合います。

• 空間計算量: \(O(N^2)\)

  • ペナルティを管理する二次元配列のサイズが \(N \times N\) となります。

実装のポイント

• インデックスの調整: 問題文では社員番号が \(1\) から \(N\) で与えられますが、プログラム内では \(0\) から \(N-1\) として扱うと配列のアクセスがスムーズになります。

• 初期値の設定: 最大値を保持する変数 max_score の初期値は、負の大きな値（例えば \(-2 \times 10^{18}\) など）にしておきます。ペナルティの合計が能力値の合計を上回り、答えが負になる可能性があるためです。

• 二重カウントの防止: ペア \((u, v)\) のペナルティを引く際、二重ループの範囲を工夫する（j を i + 1 から始めるなど）ことで、同じペアを 2 回引いてしまうミスを防げます。

  ソースコード

import sys
from itertools import combinations

def solve():
    # Read all input at once and split by whitespace for efficient parsing
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # Parse N (number of employees), M (number of bad compatibility pairs), 
    # and K (number of employees to be chosen for the team)
    N = int(input_data[0])
    M = int(input_data[1])
    K = int(input_data[2])
    
    # Parse the capability values A_i for each employee
    # A is 0-indexed, so employee i has capability A[i]
    A = [int(x) for x in input_data[3:3+N]]
    
    # penalty_matrix[u][v] stores the reduction in total capability 
    # if both employee u and employee v are in the same team.
    # Employees are converted from 1-indexed to 0-indexed.
    penalty_matrix = [[0] * N for _ in range(N)]
    idx = 3 + N
    for _ in range(M):
        u = int(input_data[idx]) - 1
        v = int(input_data[idx+1]) - 1
        b = int(input_data[idx+2])
        penalty_matrix[u][v] = b
        penalty_matrix[v][u] = b
        idx += 3
        
    # Initialize max_score with a value smaller than any possible team score.
    # The minimum possible score is roughly 1*10^9 - (18*17/2)*10^9 = -152*10^9.
    max_score = -2 * 10**18
    
    # Iterate over all possible ways to choose exactly K employees out of N.
    # itertools.combinations is an efficient iterator implemented in C.
    for combo in combinations(range(N), K):
        current_score = 0
        
        # Calculate the total capability for the chosen subset of employees.
        # We add each employee's capability and subtract penalties for every pair.
        for i, u in enumerate(combo):
            # Add individual capability
            current_score += A[u]
            
            # Retrieve the row of penalties for the current employee u
            p_u = penalty_matrix[u]
            
            # Check all other employees v in the combo that haven't been processed yet
            # to avoid double-counting pairs (u, v) and (v, u).
            for j in range(i + 1, K):
                v = combo[j]
                # Subtract the penalty if both u and v are in the team
                current_score -= p_u[v]
        
        # Update the maximum total capability found across all combinations
        if current_score > max_score:
            max_score = current_score
            
    # Output the final maximum team capability
    print(max_score)

if __name__ == '__main__':
    solve()

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