概要

棚に並んだ \(N\) 個のお菓子から、長さ \(K\) 以下の連続区間を \(M\) 個（重複なし）選び、それぞれの区間の要素の総和の合計を最大化する問題です。全区間の総和を計算し、大きい方から \(M\) 個選べばよいです。

考察

重要な気づき

各セットは「連続する区間 \([l, r]\)（長さ \(K\) 以下）」であり、異なるセット同士で区間が重なってもよく、完全に一致する区間を2回使うことだけが禁止されています。

つまり、長さ \(K\) 以下の全ての区間は候補であり、それぞれ最大1回ずつ使えます。この中から \(M\) 個を選んで満足度の合計を最大化したいので、全候補の中から満足度が大きいものを上位 \(M\) 個選べばよいという、実は貪欲法で解ける問題です。

区間の総数

長さ \(k\)（\(1 \leq k \leq K\)）の区間は \(N - k + 1\) 個あるので、候補の総数は：

\[\sum_{k=1}^{K} (N - k + 1) = NK - \frac{K(K-1)}{2}\]

\(N = K = 3000\) のとき約 \(4.5 \times 10^6\) 個となり、全列挙が十分可能な範囲です。

区間の総和の高速計算

各区間 \([l, r]\) の総和を毎回ナイーブに計算すると \(O(K)\) かかりますが、累積和（prefix sum）を前計算しておけば \(O(1)\) で求められます。

\[\text{sum}(l, r) = \text{prefix}[r+1] - \text{prefix}[l]\]

ここで \(\text{prefix}[i] = A_1 + A_2 + \cdots + A_i\)、\(\text{prefix}[0] = 0\) です。

アルゴリズム

1. 累積和の計算: 配列 \(A\) の累積和 \(\text{prefix}\) を \(O(N)\) で計算する。
2. 全区間の総和を列挙: 各始点 \(l\)（\(0 \leq l < N\)）について、長さ \(1\) から \(\min(K, N-l)\) までの区間の総和を累積和を使って \(O(1)\) で計算し、リストに格納する。
3. 降順ソート: 得られた全区間の総和を降順にソートする。
4. 上位 \(M\) 個の合計を出力: ソート結果の先頭 \(M\) 個の総和を取り、その合計を出力する。

具体例

\(N=3, M=2, K=2, A=[3, -1, 4]\) の場合：

区間	総和
\([1,1]\)	\(3\)
\([2,2]\)	\(-1\)
\([3,3]\)	\(4\)
\([1,2]\)	\(2\)
\([2,3]\)	\(3\)

降順に並べると \(4, 3, 3, 2, -1\)。上位 \(M=2\) 個を選ぶと \(4 + 3 = 7\) が答えです。

計算量

• 時間計算量: \(O(NK \log(NK))\)
  • 区間の列挙に \(O(NK)\)、ソートに \(O(NK \log(NK))\)
  • \(N = K = 3000\) のとき約 \(4.5 \times 10^6\) 個のソートで十分高速
• 空間計算量: \(O(NK)\)
  • 全区間の総和を格納するリスト分

実装のポイント

• 累積和は 0-indexed で管理すると、区間 \([l, r]\)（0-indexed）の総和が prefix[r+1] - prefix[l] とシンプルに書けます。

• おいしさが負の値を取りうるため、\(M\) 個選ばなければならない場合は負の総和の区間も選ぶ必要があります。ソートして上位 \(M\) 個を取る方法なら自然に処理されます。

• Python では約 \(4.5 \times 10^6\) 要素のリスト生成とソートは数秒程度かかる可能性があるため、sys.stdin.buffer.read() による高速入力を使うと安心です。

  ソースコード

import sys
import heapq

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    M = int(input_data[idx]); idx += 1
    K = int(input_data[idx]); idx += 1
    A = [int(input_data[idx + i]) for i in range(N)]
    idx += N

    # Prefix sum
    prefix = [0] * (N + 1)
    for i in range(N):
        prefix[i + 1] = prefix[i] + A[i]

    def range_sum(l, r):
        # sum of A[l..r] (0-indexed, inclusive)
        return prefix[r + 1] - prefix[l]

    # We need to pick M distinct intervals [l,r] with r-l+1 <= K (1-indexed, but we use 0-indexed internally)
    # to maximize sum of their range sums.
    # 
    # Total number of such intervals can be up to ~N*K which for N=3000, K=3000 is ~4.5M.
    # M can also be up to ~4.5M.
    # 
    # Strategy: compute all interval sums, sort descending, take top M.
    # But N*K can be up to 9M which might be tight but should work.
    #
    # Actually, total intervals = sum_{k=1}^{K} (N-k+1) = N*K - K*(K-1)/2
    # For N=K=3000, that's 3000*3000 - 3000*2999/2 = 9000000 - 4498500 = 4501500
    # Computing all and sorting ~4.5M elements should be feasible.
    
    # However, we need to be careful: if M is large and many sums are negative, 
    # we must still pick M intervals (the problem says M <= total number of intervals).
    # We want the top M sums.
    
    # Let's generate all interval sums efficiently.
    # For each starting position l (0-indexed), and length len from 1 to K:
    #   r = l + len - 1, must have r < N
    #   sum = prefix[r+1] - prefix[l]
    
    # To avoid memory issues with very large lists, let's use a heap-based approach
    # or just generate and sort.
    
    # For N=3000, K=3000: ~4.5M values. Generating a list and sorting should work.
    
    total_intervals = 0
    for k in range(1, K + 1):
        total_intervals += (N - k + 1)
    
    # Generate all sums
    sums = []
    for l in range(N):
        max_len = min(K, N - l)
        for length in range(1, max_len + 1):
            r = l + length - 1
            s = prefix[r + 1] - prefix[l]
            sums.append(s)
    
    # Sort descending
    sums.sort(reverse=True)
    
    # Take top M
    result = 0
    for i in range(M):
        result += sums[i]
    
    print(result)

main()

---

この解説は claude4.6opus-thinking によって生成されました。