Overview

Consider all possible sums (satisfaction values) of contiguous subarrays of length at most \(K\), and maximize the total when selecting the top \(M\) largest ones (i.e., find the sum of the top \(M\) values among all “length-restricted subarray sums”).

Analysis

Key Observations

• Each set is a contiguous interval \([l,r]\), so the satisfaction is the subarray sum \(\sum_{i=l}^{r} A_i\).
• Intervals may overlap, and the only restriction is “the same interval cannot be chosen twice,” so the optimal strategy is simply to take the \(M\) largest subarray sums among all subarrays of length \(\le K\) and add them up.

Why a Naive Solution Is Too Slow

• The number of intervals is \(\sum_{k=1}^{K}(N-k+1)=O(NK)\), which is at most about \(4.5\) million (when \(N=K=3000\)).
• Enumerating all of them, storing them in an array, sorting, and taking the top \(M\):
  • In Python, storing all subarray sums is memory-heavy
  • Sorting is \(O(NK\log(NK))\), which can be large

This makes the implementation burden significant.

Solution Strategy

Instead of directly extracting the “top \(M\),” we enable fast computation of: - How many subarray sums are \(\ge x\) (count) - What is the total of those subarray sums (sum)

Then we binary search on \(x\) to find the “threshold” (the \(M\)-th largest value).

Algorithm

1. Represent Subarray Sums Using Prefix Sums

Define the prefix sum \(S\) as: - \(S[0]=0\) - \(S[i]=A_1+\cdots+A_i\)

Then the sum of interval \([l,r]\) is [ \text{sum}(l,r)=S[r]-S[l-1] ]

The length restriction \(r-l+1\le K\) is equivalent to [ (l-1) \ge r-K ] That is, when the right endpoint \(r\) is fixed, the valid prefix sum indices \(p=l-1\) satisfy [ p \in [r-K,\, r-1] ]

2. Count and Sum Subarray Sums \(\ge x\) (calc_ge(x))

The condition [ S[r]-S[p]\ge x ] is equivalent to [ S[p]\le S[r]-x ]

So for each \(r=1..N\): - Maintain the values \(S[p]\) for \(p\in[r-K, r-1]\) in a data structure (sliding window) - Count how many \(S[p]\) satisfy \(S[p] \le \text{bound}=S[r]-x\) (count \(c\)) - Compute the sum of those \(S[p]\) values (\(sp\))

Then the total of “qualifying subarray sums” with right endpoint \(r\) is [ \sum (S[r]-S[p]) = c\cdot S[r] - sp ] Summing over all \(r\) gives us: - The total count of qualifying intervals: total_c - The total sum of qualifying subarray sums: total_s

Data Structure: Fenwick Tree (BIT)

Since \(S[p]\) can be negative or very large, it cannot be used directly as a BIT index.
So we apply coordinate compression (vals = sorted(set(S))) and maintain on the BIT: - bitc: counts - bits: sum of values

Furthermore, to enforce the length restriction \(\le K\), as \(r\) increases by 1: - Add the newly usable \(S[r-1]\) - Remove the expired \(S[r-K-1]\) (when \(r-K-1\ge 0\))

This maintains the sliding window (corresponding to add(..., +1) and add(..., -1) in the code).

3. Determine the Threshold via Binary Search

Using calc_ge(x) to obtain “the count of subarray sums \(\ge x\)”: - Binary search for the maximum \(x\) such that the count \(\ge M\). Call this threshold.

threshold is the boundary value determining whether a subarray sum is among the top \(M\) (there may be multiple subarray sums equal to this value).

4. Reconstruct the Sum of the Top \(M\)

• Let \(c_{gt}\) and \(s_{gt}\) be the count and sum, respectively, of subarray sums \(\ge\) threshold+1 (i.e., strictly greater than threshold), obtained via calc_ge(threshold+1).
• Among the top \(M\), all subarray sums strictly greater than threshold are definitely included, contributing \(s_{gt}\).
• The remaining \(M-c_{gt}\) slots are filled with subarray sums exactly equal to threshold, so: [ \text{answer} = s{gt} + (M-c{gt})\cdot \text{threshold} ]

Complexity

• Time complexity:
  calc_ge(x) runs in \(O(N\log N)\), and binary search performs about \(O(\log(\text{value range}))\approx 60\) iterations, so
  [ O(N\log N\log W) ] (where \(W\) is the range of possible subarray sums; in the implementation, we search over \([-10^{18},10^{18}]\))
• Space complexity:
  Coordinate compression and BIT use [ O(N) ]

Implementation Notes

• Coordinate compression: Compress the prefix sums \(S\) via sorted(set(S)) to use as BIT indices.
• Maintain two BITs:
  
  • bitc for counts
  • bits for sum of values
    By updating and querying both simultaneously, the total \(\sum(S[r]-S[p])\) can be efficiently computed as \(c\cdot S[r]-sp\).
• Sliding window management for the length restriction \(K\): Since the valid range of \(p\) changes for each \(r\):
  • Add: \(S[r-1]\)
  • Remove: \(S[r-K-1]\) (if it exists)

Make sure not to forget these operations. - Final sum construction (handling ties):
When there are multiple subarray sums exactly equal to threshold, first take all values strictly greater than threshold, then fill the remainder with threshold. This ensures correctness.

Source Code

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()

---

This editorial was generated by gpt-5.2-high.