概要

\(N\) 人の社員の中から \(K\) 人を選ぶチーム編成において、各社員の能力値の合計から、同じチームに入った相性の悪いペアによる減点を引いた「チームの総合力」を最大化する問題。

考察

この問題は、社員の選び方によってチームの総合力が変わる組み合わせ最適化問題です。
愚直に考えると、\(N\) 人の中から \(K\) 人を選ぶすべての組み合わせを試せばよいですが、これは二項係数 \({}_N \mathrm{C}_K\) 通りあります。制約より \(N \leq 18\) なので、最大でも \({}_{18} \mathrm{C}_9 = 48620\) 程度であり、現実的な計算量で全探索が可能です。

一方で、それぞれの組み合わせに対して、チームに含まれる「相性の悪いペア」の減点を計算する必要があります。
これは、事前にペア情報を管理しておくことで効率よく計算できます。また、ペア \((u, v)\) は \(u < v\) の順で保持しておくと、扱いやすくなります。

さらに、減点は「選んだペアがチームに両方含まれている場合」にのみ発生するので、各チーム候補に対して、含まれるペアを逐一確認すればOKです。

このように、全組み合わせを試しながら、各組み合わせでの得点を高速に求めることで、問題を解くことができます。

アルゴリズム

1. 社員の能力値 \(A_i\) と、相性の悪いペア情報 \((U_j, V_j, B_j)\) を読み込む。
2. 相性ペアを \((u, v)\) をキー、減点 \(b\) を値とする辞書 bad に保存（\(u < v\) になるよう調整）。
3. itertools.combinations を使い、\(N\) 人から \(K\) 人を選ぶすべての組み合わせを生成。
4. 各組み合わせ（チーム候補）について：
   • まずチームメンバーの能力値の合計を計算。
   • 次に、そのチームに含まれる相性の悪いペアがあれば、その減点を引く。
5. すべてのチーム候補の中で最大の総合力を記録し、出力。

計算量

• 時間計算量: \(O\left(\binom{N}{K} \cdot M\right)\)
  • 組み合わせの数は最大で \({}_{18}\mathrm{C}_9 \approx 4.8 \times 10^4\)
  • 各組み合わせに対し、最大 \(M\) 個のペアを確認する必要がある。
• 空間計算量: \(O(M + N)\)
  • 能力値リストとペア情報を格納するのに必要な領域。

実装のポイント

• 社員番号は入力で 1-indexed だが、内部処理では 0-indexed にしておくと扱いやすい。

• ペア \((u, v)\) は必ず \(u < v\) の順で保存することで、辞書引きが簡単になる。

• set を使ってチームメンバーに特定の社員が含まれるかどうかを高速に判定している（リストだと \(O(K)\)、set なら平均 \(O(1)\)）。

  ソースコード

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    idx = 0
    N = int(data[idx]); idx += 1
    M = int(data[idx]); idx += 1
    K = int(data[idx]); idx += 1
    
    A = [int(data[idx + i]) for i in range(N)]
    idx += N
    
    # 相性の悪いペアを格納する辞書: (u, v) -> cost
    bad = {}
    for _ in range(M):
        u = int(data[idx]) - 1; idx += 1
        v = int(data[idx]) - 1; idx += 1
        b = int(data[idx]); idx += 1
        if u > v:
            u, v = v, u
        bad[(u, v)] = b

    # 全てのK人組み合わせを試す
    from itertools import combinations
    max_power = -float('inf')
    
    for team in combinations(range(N), K):
        power = sum(A[i] for i in team)
        # チーム内の相性ペアをチェック
        team_set = set(team)
        for (u, v), cost in bad.items():
            if u in team_set and v in team_set:
                power -= cost
        if power > max_power:
            max_power = power
            
    print(max_power)

if __name__ == "__main__":
    main()

この解説は qwen3-coder-480b によって生成されました。