Official

B - お菓子の詰め合わせ / Assortment of Sweets Editorial by admin

DeepSeek V3

Overview

This is a problem where you select M intervals of consecutive sweets (each of length at most K) to maximize the total satisfaction.

Analysis

From the problem conditions, we can observe the following: - Each set consists of at least 1 and at most K consecutive sweets - The same interval cannot be selected more than once, but overlapping intervals can be selected - The satisfaction is the sum of deliciousness values within the interval - The total number of possible intervals is at most about \(N \times K\) (under the constraints, at most 3000×3000 = 9,000,000)

A straightforward approach is to compute the satisfaction of all possible intervals and select the top M in descending order. Under the constraints, the number of intervals is at most 9,000,000, and sorting this may seem time-consuming at first glance, but in practice, when \(N=3000\), the number of intervals is about 4,500,000, which is within a sortable range even in Python.

Algorithm

  1. Prefix sum computation: Precompute the prefix sum from the deliciousness array. This allows computing the sum of any interval in \(O(1)\).
  2. Compute satisfaction for all intervals: For all intervals of length 1 to K, compute the satisfaction (interval sum) and store them in a list.
  3. Sort and sum: Sort the satisfaction list in descending order and output the sum of the top M values.

Complexity

  • Time complexity: \(O(NK + NK \log(NK))\)
    • The number of intervals is approximately \(\frac{K(2N-K+1)}{2}\) ≒ \(O(NK)\)
    • Sorting takes \(O(NK \log(NK))\)
  • Space complexity: \(O(NK)\)
    • To store the satisfaction values of all intervals

Implementation Notes

  • Use prefix sums to speed up interval sum computation

  • Enumerate all valid intervals with a nested loop

  • Specify reverse=True during sorting for descending order

  • Simply compute the sum of the top M values

    Source Code

def main():
    import sys
    data = sys.stdin.read().split()
    it = iter(data)
    N = int(next(it)); M = int(next(it)); K = int(next(it))
    A = list(map(int, [next(it) for _ in range(N)]))
    
    prefix = [0] * (N+1)
    for i in range(1, N+1):
        prefix[i] = prefix[i-1] + A[i-1]
    
    S = []
    for length in range(1, K+1):
        for start in range(0, N - length + 1):
            end = start + length - 1
            total = prefix[end+1] - prefix[start]
            S.append(total)
    
    S.sort(reverse=True)
    ans = sum(S[:M])
    print(ans)

if __name__ == "__main__":
    main()

This editorial was generated by deepseekv3.

posted:
last update: