[Python] 半分全列挙と二分探索 JOI2008本選C

JOI2008本選C

問題

N個の点数Pが書かれた的でダーツをする。はじめに謎の点数Mが設定され、その点数を超えないよう矢を的に向かって投げる。Mを超えると０点になる。矢は４本あるが全てを投げきらなくともよい。合計点を最大化したい。

制約

$1 \leqq N \leqq 1000$
$1 \leqq p_i \leqq 10^8$

考え方

４本投げるので単純探索を4重ループを書くと$O(N^4)$で間に合いません。「4 本の矢を投げる」ことを，「まず 2 本の矢をまとめて投げて，次に 2 本の矢をまとめて投げる」と考えます。実装としては、まず2 本の点数を半分全列挙 $O(N^2)$する。その結果を使い点数制限内で最大値を二分探索 $O(\log N)$することで、制限時間内に全探索 $O(N^2\log N)$することができそうです。

import bisect

N, M = map(int, input().split())
P = [0]
for n in range(N):
    P.append(int(input()))

# 2本の点数を半分全列挙
S = []
for p1 in P:
    for p2 in P:
        S.append(p1+p2)
# ソートして二分探索に備える
S.sort()

ans = 0
for s in S:
    if M < s:
        break
    # 点数制限内での最大値を二分探索
    i = bisect.bisect_right(S, M-s)-1
    # それがこれまでの最大なら答えを更新
    ans = max(s+S[i], ans)

print(ans)