Overview

Given \(N\) sweets lined up on a shelf, we select \(M\) contiguous intervals of length at most \(K\) (without duplicates), and maximize the total sum of elements in each interval. We compute the sum of all possible intervals and choose the top \(M\) largest ones.

Analysis

Key Insight

Each set is a “contiguous interval \([l, r]\) (of length at most \(K\))”, and different sets may overlap with each other — only using the exact same interval twice is forbidden.

In other words, all intervals of length at most \(K\) are candidates, and each can be used at most once. Since we want to select \(M\) of these to maximize the total satisfaction, we simply pick the top \(M\) candidates with the largest sums. This is actually a problem solvable by a greedy approach.

Total Number of Intervals

The number of intervals of length \(k\) (\(1 \leq k \leq K\)) is \(N - k + 1\), so the total number of candidates is:

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

When \(N = K = 3000\), this is approximately \(4.5 \times 10^6\), which is well within the range for full enumeration.

Fast Computation of Interval Sums

Naively computing the sum of each interval \([l, r]\) takes \(O(K)\), but by precomputing a prefix sum, we can compute it in \(O(1)\).

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

Here, \(\text{prefix}[i] = A_1 + A_2 + \cdots + A_i\) and \(\text{prefix}[0] = 0\).

Algorithm

1. Compute prefix sums: Compute the prefix sum array \(\text{prefix}\) of array \(A\) in \(O(N)\).
2. Enumerate all interval sums: For each starting point \(l\) (\(0 \leq l < N\)), compute the sum of intervals of length \(1\) to \(\min(K, N-l)\) in \(O(1)\) using the prefix sum, and store them in a list.
3. Sort in descending order: Sort all the computed interval sums in descending order.
4. Output the sum of the top \(M\): Take the first \(M\) values from the sorted result and output their total sum.

Concrete Example

For \(N=3, M=2, K=2, A=[3, -1, 4]\):

Interval	Sum
\([1,1]\)	\(3\)
\([2,2]\)	\(-1\)
\([3,3]\)	\(4\)
\([1,2]\)	\(2\)
\([2,3]\)	\(3\)

Sorted in descending order: \(4, 3, 3, 2, -1\). Selecting the top \(M=2\) gives \(4 + 3 = 7\) as the answer.

Complexity

• Time complexity: \(O(NK \log(NK))\)
  • Enumerating intervals takes \(O(NK)\), sorting takes \(O(NK \log(NK))\)
  • When \(N = K = 3000\), sorting approximately \(4.5 \times 10^6\) elements is sufficiently fast
• Space complexity: \(O(NK)\)
  • For the list storing all interval sums

Implementation Notes

• Managing prefix sums with 0-indexed allows the sum of interval \([l, r]\) (0-indexed) to be written simply as prefix[r+1] - prefix[l].

• Since tastiness values can be negative, when we are required to select exactly \(M\) intervals, we may need to include intervals with negative sums. The approach of sorting and taking the top \(M\) handles this naturally.

• In Python, generating and sorting a list of approximately \(4.5 \times 10^6\) elements may take several seconds, so it is advisable to use fast input with sys.stdin.buffer.read().

  Source Code

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

---

This editorial was generated by claude4.6opus-thinking.