概要

\(N\) 個の要素からなる数列から、長さ \(K\) 以下の連続する部分区間を \(M\) 個選び、それらの和の合計を最大化する問題です。同じ区間を重複して選ぶことはできませんが、異なる区間同士が重なることは許容されています。

考察

この問題のポイントは、「選ぶ区間同士が重なってもよい」という点です。

もし「重なってはいけない」という制約があれば動的計画法（DP）などが必要になりますが、重なってもよい（ただし区間 \([l, r]\) としては別のものでなければならない）ため、単純に「あり得るすべての区間の中から、和が大きい順に \(M\) 個選ぶ」という貪欲法が成立します。

1. 候補となる区間の数

選ぶことができる区間の条件は「長さが \(1\) 以上 \(K\) 以下」です。 - 長さ 1 の区間：\(N\) 個 - 長さ 2 の区間：\(N-1\) 個 - … - 長さ \(K\) の区間：\(N-K+1\) 個

これらをすべて合わせると、最大で \(N \times K\) 個程度の区間が存在します。制約の \(N=3000\) を代入すると、区間の総数は最大で約 450 万個（\(\frac{3000 \times 3001}{2} \approx 4.5 \times 10^6\)）となります。

2. 区間和の計算

すべての区間の和を効率よく計算するために、累積和を利用します。 数列 \(A\) の累積和 \(P\) を \(P_i = A_1 + \dots + A_i\) と定義すると、区間 \([l, r]\) の和は \(P_r - P_{l-1}\) で計算できます。これにより、各区間の和を \(O(1)\) で求めることができます。

アルゴリズム

1. 累積和の準備: 数列 \(A\) の累積和 \(P\) を計算します。
2. 全区間和の列挙: 長さ \(1\) から \(K\) までのすべての連続する区間について、その和を計算してリストに格納します。
3. ソート: 格納した区間和を降順（大きい順）にソートします。
4. 合計: ソートしたリストの先頭から \(M\) 個の要素を足し合わせたものが答えとなります。

計算量

• 時間計算量: \(O(NK \log(NK))\)
  • 全区間の列挙に \(O(NK)\) かかります。
  • リストのソートに \(O(NK \log(NK))\) かかります。\(NK \approx 4.5 \times 10^6\) の場合、Python の高速なソートアルゴリズム（Timsort）であれば制限時間内に処理可能です。
• 空間計算量: \(O(NK)\)
  • すべての区間和を保持するために \(O(NK)\) のメモリを使用します。

実装のポイント

• 効率的なリスト作成: Python では for ループで 1 つずつ append するよりも、map やスライス、extend を活用することで、大量のデータ処理を高速化できます。

• メモリ管理: \(4.5 \times 10^6\) 個の整数を保持するリストはメモリを消費します。不要なデータは早めに削除したり、スライスで必要な範囲だけを合計したりする工夫が有効です。

• 累積和の活用: P[r] - P[l] の計算を operator.sub と map で一括処理することで、内部ループを高速に回しています。

  ソースコード

import sys
import operator
import itertools

def solve():
    """
    高橋君のお菓子の詰め合わせ問題を解く。
    N個のお菓子から長さK以下の連続する区間をM個選び、その和の合計を最大化する。
    """
    # 標準入力からすべてのデータを読み込み、スペースで分割する
    # N=3000程度なので、入力サイズは大きくなく、一括読み込みが効率的
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    it = iter(input_data)
    try:
        # N: お菓子の数, M: 作るセットの数, K: 1セットのお菓子の最大個数
        N = int(next(it))
        M = int(next(it))
        K = int(next(it))
        # 各お菓子のおいしさ
        A = list(map(int, it))
    except StopIteration:
        return
    
    # 累積和を計算する: P[i] は A[0]...A[i-1] の和
    # これにより、区間 [l, r] の和を P[r] - P[l-1] で O(1) で計算できる
    P = [0] * (N + 1)
    s = 0
    for i in range(N):
        s += A[i]
        P[i+1] = s
    
    # すべての可能な区間の和を格納するリスト
    all_sums = []
    # 速度向上のためにメソッドをローカル変数にキャッシュ
    all_sums_extend = all_sums.extend
    sub = operator.sub
    
    # 長さ 1 から K までのすべての連続する区間の和を計算する
    # 各長さについて、スライシングと map を組み合わせて効率的に計算
    for length in range(1, K + 1):
        # P[length:] は P[length], P[length+1], ..., P[N]
        # P は P[0], P[1], ..., P[N]
        # map(sub, P[length:], P) は P[length]-P[0], P[length+1]-P[1], ..., P[N]-P[N-length] を生成する
        # これは長さ 'length' のすべての区間の和に対応する
        all_sums_extend(map(sub, P[length:], P))
    
    # 区間の和を降順（大きい順）にソートする
    # Pythonの sort (Timsort) は非常に高速であり、4.5e6 要素程度なら1秒以内で処理可能
    all_sums.sort(reverse=True)
    
    # 最大の満足度を得るために、上位 M 個の区間の和を選択する
    # メモリ節約のため、M個より多い分は削除してから合計を計算する
    if len(all_sums) > M:
        del all_sums[M:]
    
    # 合計値を出力
    print(sum(all_sums))

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-thinking によって生成されました。