Official
B - お菓子の詰め合わせ / Assortment of Sweets Editorial by admin
DeepSeek V3概要
連続するお菓子の区間(長さK以下)からM個選び、満足度の合計を最大化する問題です。
考察
問題の条件から、以下のことが分かります: - 各セットは連続する1個以上K個以下のお菓子から構成される - 同じ区間を2回以上選べないが、重なる区間は選べる - 満足度は区間内のおいしさの総和 - 選べる区間の総数は最大で約 \(N \times K\) 個(制約では最大3000×3000=9,000,000個)
素朴なアプローチとして、すべての可能な区間の満足度を計算し、大きい順にM個選ぶ方法が考えられます。制約上、区間の数は最大9,000,000個となり、これをソートするのは一見時間がかかりそうに見えますが、実際には \(N=3000\) のときの区間数は約4,500,000個で、Pythonでもソート可能な範囲内です。
アルゴリズム
- 累積和の計算: おいしさの配列から累積和を事前計算します。これにより任意の区間の和を \(O(1)\) で求められます。
- 全区間の満足度計算: 長さ1からKまでのすべての区間について、満足度(区間和)を計算しリストに格納します。
- ソートと合計: 満足度リストを降順にソートし、上位M個の合計を出力します。
計算量
- 時間計算量: \(O(NK + NK \log(NK))\)
- 区間の数は約 \(\frac{K(2N-K+1)}{2}\) ≒ \(O(NK)\)
- ソートに \(O(NK \log(NK))\)
- 空間計算量: \(O(NK)\)
- すべての区間の満足度を格納するため
実装のポイント
累積和を使うことで区間和の計算を高速化
二重ループで全ての有効な区間を列挙
ソート時に
reverse=Trueを指定して降順ソート上位M個の合計を単純に計算
ソースコード
def main():
import sys
data = sys.stdin.read().split()
it = iter(data)
N = int(next(it)); M = int(next(it)); K = int(next(it))
A = list(map(int, [next(it) for _ in range(N)]))
prefix = [0] * (N+1)
for i in range(1, N+1):
prefix[i] = prefix[i-1] + A[i-1]
S = []
for length in range(1, K+1):
for start in range(0, N - length + 1):
end = start + length - 1
total = prefix[end+1] - prefix[start]
S.append(total)
S.sort(reverse=True)
ans = sum(S[:M])
print(ans)
if __name__ == "__main__":
main()
この解説は deepseekv3 によって生成されました。