D - チームビルディング / Team Building 解説 by admin
GPT 5.2 High概要
\(N \le 18\) なので、選ぶ社員集合をビットマスクで全探索しつつ、「能力の合計 − 相性ペナルティ」の値を部分集合DPで高速に求め、ちょうど \(K\) 人のときの最大値を出します。
考察
社員集合 \(S\) を選んだときの総合力は
- 能力の合計:\(\sum_{i \in S} A_i\)
- ペナルティ:\(S\) 内に含まれる相性の悪いペア \((u,v)\) ごとに \(B\) を引く
なので、素朴には各集合 \(S\) について 1. \(\sum A_i\) を計算 2. すべてのペアを調べてペナルティを足す(あるいは \(M\) 個の相性ペアを全部チェックする)
とやりたくなります。
しかし、全ての部分集合は \(2^N\) 個あり、各集合ごとにペナルティを愚直に数えると(例えば \(O(M)\) かかると) - \(2^N \cdot M\) となり、\(N=18\) でも Python では重くなりがちです。
そこで重要な気づきは:
- 部分集合を 1 要素ずつ増やす形でDPすると、増えた分(新しく追加した社員と、既に入っている社員との相性ペナルティ)だけを引けばよい
- 「ある社員 \(i\) と、集合 \(S\) にいる全員とのペナルティ合計」\(\sum_{j \in S} w[i][j]\) を高速に取り出せれば、1遷移が \(O(1)\) になる
という点です。
アルゴリズム
1. ビットマスクで集合を表す
長さ \(N\) のビット列 mask を使い、ビット \(i\) が 1 なら社員 \(i\) を選んでいる、と表します(0-indexed)。
2. 前計算:penSum[i][mask]
penSum[i][mask] = \sum_{j \in mask} w[i][j] を全ての \(i\) と mask について作ります。
これは部分集合DPで計算できます。mask の最下位の 1 ビットを lsb、その位置を \(j\) とすると
penSum[i][mask] = penSum[i][mask ^ lsb] + w[i][j]
で求まります。これにより penSum が全体で \(O(N \cdot 2^N)\) で作れます。
3. 部分集合DP:score[mask]
score[mask] を「集合 mask を選んだときの総合力」とします。
mask の最下位の 1 ビットが社員 \(i\) を表すとき、prev = mask ^ lsb は「社員 \(i\) を除いた集合」です。
このとき prev から mask に増えた分は:
- 能力:\(+A_i\)
- 新しく発生するペナルティ:\(-\sum_{j \in prev} w[i][j]\)
よって
[ score[mask] = score[prev] + A_i - penSum[i][prev] ]
となります。
この式だと、相性ペア \((u,v)\) のペナルティは「後から追加された側」でちょうど 1 回だけ引かれるため、二重に引くことも漏れることもありません。
4. 答え
全ての mask について、bit_count()(集合サイズ)が \(K\) のものだけ見て score[mask] の最大値を取ります。
計算量
- 時間計算量: \(O(N \cdot 2^N)\)
(penSumの構築が \(O(N \cdot 2^N)\)、scoreのDPも \(O(2^N)\)) - 空間計算量: \(O(N \cdot 2^N)\)
(penSumが支配的。scoreは \(O(2^N)\))
実装のポイント
最下位の1ビット(LSB)を使う
lsb = mask & -mask、i = lsb.bit_length() - 1で、1要素追加の形に分解できます。ペナルティの加算を高速化する前計算が肝
penSum[i][prev]を \(O(1)\) で取れるので、scoreの遷移が高速になります。Pythonでは配列を使って高速化
コードではarray('q')(64bit整数)を使い、リストよりメモリ・速度面で有利にしています(答えは最大で \(10^9\) オーダーが足し引きされるので 64bit で安全)。ソースコード
import sys
from array import array
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
it = iter(data)
N = next(it)
M = next(it)
K = next(it)
A = [next(it) for _ in range(N)]
w = [[0] * N for _ in range(N)]
for _ in range(M):
u = next(it) - 1
v = next(it) - 1
b = next(it)
w[u][v] = b
w[v][u] = b
size = 1 << N
penSum = []
for i in range(N):
arr = array('q', [0]) * size
for mask in range(1, size):
lsb = mask & -mask
j = lsb.bit_length() - 1
arr[mask] = arr[mask ^ lsb] + w[i][j]
penSum.append(arr)
score = array('q', [0]) * size
for mask in range(1, size):
lsb = mask & -mask
i = lsb.bit_length() - 1
prev = mask ^ lsb
score[mask] = score[prev] + A[i] - penSum[i][prev]
ans = -10**30
for mask in range(size):
if mask.bit_count() == K:
v = score[mask]
if v > ans:
ans = v
print(ans)
if __name__ == "__main__":
main()
この解説は gpt-5.2-high によって生成されました。
投稿日時:
最終更新: