概要

長さ \(K\) 以下の連続区間の総和（満足度）をすべて考え、その中から大きいものを上位 \(M\) 個選んだときの合計値を最大化します（＝「長さ制限付き部分列和の上位 \(M\) 個の総和」を求める問題です）。

考察

重要な気づき

• 1つのセットは区間 \([l,r]\)（連続）なので、満足度は区間和 \(\sum_{i=l}^{r} A_i\) です。
• 区間が重なっても良く、「同じ区間を2回選べない」だけなので、結局
  長さ \(\le K\) のすべての区間和の中から、大きいものを \(M\) 個取って足すのが最適です。

素朴解が厳しい理由

• 区間数は \(\sum_{k=1}^{K}(N-k+1)=O(NK)\) で、最大で約 \(4.5\) 百万個（\(N=K=3000\)）。
• これらを全部列挙して配列に入れ、ソートして上位 \(M\) 個を取ると、
  • Pythonでは「区間和を全部保持」がメモリ的に重い
  • ソートも \(O(NK\log(NK))\) と大きくなりがち
    で、実装上の負担が大きいです。

解決方針

「上位 \(M\) 個」を直接取り出す代わりに、 - ある値 \(x\) 以上の区間和が 何個あるか（個数） - それらの区間和の 合計がいくつか（合計値） を高速に数えられるようにし、\(x\) を二分探索して「\(M\) 番目に大きい値（閾値）」を見つけます。

アルゴリズム

1. 累積和で区間和を表す

累積和 \(S\) を - \(S[0]=0\) - \(S[i]=A_1+\cdots+A_i\) とすると、区間 \([l,r]\) の和は [ \text{sum}(l,r)=S[r]-S[l-1] ] です。

長さ制限 \(r-l+1\le K\) は [ (l-1) \ge r-K ] つまり、右端 \(r\) を固定したとき、使える左側の累積和の添字 \(p=l-1\) は [ p \in [r-K,\, r-1] ] の範囲になります。

2. 「区間和 \(\ge x\)」の個数と合計を数える（calc_ge(x)）

条件 [ S[r]-S[p]\ge x ] は [ S[p]\le S[r]-x ] と同値です。

そこで、各 \(r=1..N\) について - 範囲 \(p\in[r-K, r-1]\) の \(S[p]\) をデータ構造に入れておく（スライド窓） - 上限 \(\text{bound}=S[r]-x\) 以下の \(S[p]\) が何個あるか（個数 \(c\)） - それらの \(S[p]\) の合計（\(sp\)） を求めます。

このとき、\(r\) を右端とする「条件を満たす区間和」の合計は [ \sum (S[r]-S[p]) = c\cdot S[r] - sp ] となるので、全 \(r\) で足し合わせれば - 条件を満たす区間の個数 total_c - 条件を満たす区間和の総和 total_s が得られます。

データ構造：Fenwick Tree（BIT）

\(S[p]\) は負にも大きな値にもなるので、そのままBITの添字にできません。
そこで 座標圧縮（vals = sorted(set(S))）をして、BIT上で - bitc：個数（count） - bits：値の合計（sum of prefix values） を管理します。

さらに「長さ \(\le K\)」のために、\(r\) が1増えるごとに - 新しく使える \(S[r-1]\) を追加 - 古くなった \(S[r-K-1]\) を削除（\(r-K-1\ge 0\) のとき） としてスライド窓を保ちます（コードの add(..., +1) と add(..., -1)）。

3. 二分探索で閾値を決める

calc_ge(x) で「区間和 \(\ge x\) の個数」が分かるので、 - 個数 \(\ge M\) となる最大の \(x\) を二分探索で求めます。これを threshold とします。

threshold は「上位 \(M\) 個に入るかどうかの境界値」になります（同値が複数ある可能性あり）。

4. 上位 \(M\) 個の合計を復元する

• threshold+1 以上（つまり threshold より大きい）区間和の
  • 個数を \(c_{gt}\)
  • 合計を \(s_{gt}\) とします（calc_ge(threshold+1)）。
• 上位 \(M\) 個のうち、まず「thresholdより大きいもの」は全部確定で入るので合計は \(s_{gt}\)。
• 残り \(M-c_{gt}\) 個は、値がちょうど threshold のものから取ればよいので、 [ \text{answer} = s{gt} + (M-c{gt})\cdot \text{threshold} ] で求まります。

計算量

• 時間計算量:
  calc_ge(x) が \(O(N\log N)\)、二分探索が約 \(O(\log(\text{値域}))\approx 60\) 回なので
  [ O(N\log N\log W) ] （\(W\) は区間和の取りうる値域。実装では \([-10^{18},10^{18}]\) で探索）
• 空間計算量:
  座標圧縮とBITで [ O(N) ]

実装のポイント

• 座標圧縮：累積和 \(S\) を sorted(set(S)) にして BIT の添字にします。

• BITを2本持つ：
  

  • 個数用 bitc
  • 値の合計用 bits
    を同時に更新・取得すると、合計 \(\sum(S[r]-S[p])\) を \(c\cdot S[r]-sp\) で高速に計算できます。

• 長さ制限 \(K\) のスライド窓管理：各 \(r\) で「使える \(p\) の範囲」が変わるため、

  • 追加：\(S[r-1]\)
  • 削除：\(S[r-K-1]\)（存在する場合） を忘れないようにします。

• 最終合計の作り方（同値処理）：
  threshold ちょうどの区間和が複数あるときに備え、
  「threshold より大きいものを先に全部取って、残りを threshold で埋める」形にすると確実です。

  ソースコード

import sys
from bisect import bisect_right, bisect_left

def main():
    it = iter(map(int, sys.stdin.buffer.read().split()))
    N = next(it)
    M = next(it)
    K = next(it)
    A = [next(it) for _ in range(N)]

    S = [0] * (N + 1)
    for i in range(N):
        S[i + 1] = S[i] + A[i]

    vals = sorted(set(S))
    V = len(vals)
    idx_pref = [bisect_left(vals, x) + 1 for x in S]  # 1-indexed for BIT

    def calc_ge(x: int):
        bitc = [0] * (V + 1)
        bits = [0] * (V + 1)

        def add(i: int, dc: int, ds: int):
            while i <= V:
                bitc[i] += dc
                bits[i] += ds
                i += i & -i

        def query(i: int):
            c = 0
            s = 0
            while i > 0:
                c += bitc[i]
                s += bits[i]
                i -= i & -i
            return c, s

        total_c = 0
        total_s = 0
        for r in range(1, N + 1):
            add(idx_pref[r - 1], 1, S[r - 1])
            if r - K - 1 >= 0:
                add(idx_pref[r - K - 1], -1, -S[r - K - 1])

            bound = S[r] - x
            qi = bisect_right(vals, bound)
            c, sp = query(qi)
            total_c += c
            total_s += c * S[r] - sp

        return total_c, total_s

    lo = -10**18  # count >= M
    hi = 10**18   # count < M
    while hi - lo > 1:
        mid = (lo + hi) // 2
        c, _ = calc_ge(mid)
        if c >= M:
            lo = mid
        else:
            hi = mid

    threshold = lo
    c_gt, s_gt = calc_ge(threshold + 1)
    ans = s_gt + (M - c_gt) * threshold
    print(ans)

if __name__ == "__main__":
    main()

---

この解説は gpt-5.2-high によって生成されました。