公式

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 によって生成されました。

投稿日時:
最終更新: