Official

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

Gemini 3.0 Flash (Thinking)

Overview

Given a sequence of \(N\) elements, select \(M\) contiguous subintervals of length at most \(K\), and maximize the total sum of those intervals. The same interval cannot be selected more than once, but different intervals are allowed to overlap with each other.

Analysis

The key point of this problem is that “selected intervals are allowed to overlap.”

If there were a constraint that intervals must not overlap, dynamic programming (DP) or similar techniques would be needed. However, since overlapping is allowed (as long as the intervals \([l, r]\) are distinct), a simple greedy approach of “selecting the \(M\) intervals with the largest sums from all possible intervals” works.

1. Number of Candidate Intervals

The condition for selectable intervals is “length between \(1\) and \(K\) inclusive.” - Intervals of length 1: \(N\) - Intervals of length 2: \(N-1\) - … - Intervals of length \(K\): \(N-K+1\)

Combining all of these, there are at most about \(N \times K\) intervals. Substituting the constraint \(N=3000\), the total number of intervals is at most approximately 4.5 million (\(\frac{3000 \times 3001}{2} \approx 4.5 \times 10^6\)).

2. Computing Interval Sums

To efficiently compute the sum of all intervals, we use prefix sums. Define the prefix sum \(P\) of sequence \(A\) as \(P_i = A_1 + \dots + A_i\). Then the sum of interval \([l, r]\) can be computed as \(P_r - P_{l-1}\). This allows us to compute each interval sum in \(O(1)\).

Algorithm

  1. Prepare prefix sums: Compute the prefix sum array \(P\) of sequence \(A\).
  2. Enumerate all interval sums: For all contiguous intervals of length \(1\) to \(K\), compute their sums and store them in a list.
  3. Sort: Sort the stored interval sums in descending order (largest first).
  4. Sum up: The answer is the sum of the first \(M\) elements from the sorted list.

Complexity

  • Time complexity: \(O(NK \log(NK))\)
    • Enumerating all intervals takes \(O(NK)\).
    • Sorting the list takes \(O(NK \log(NK))\). When \(NK \approx 4.5 \times 10^6\), Python’s fast sorting algorithm (Timsort) can handle this within the time limit.
  • Space complexity: \(O(NK)\)
    • \(O(NK)\) memory is used to store all interval sums.

Implementation Notes

  • Efficient list construction: In Python, rather than using a for loop with append one by one, leveraging map, slicing, and extend can speed up processing of large amounts of data.

  • Memory management: A list holding \(4.5 \times 10^6\) integers consumes significant memory. It is effective to delete unnecessary data early or use slicing to sum only the required range.

  • Leveraging prefix sums: By batch-processing the computation of P[r] - P[l] using operator.sub and map, the inner loop runs efficiently.

    Source Code

import sys
import operator
import itertools

def solve():
    """
    高橋君のお菓子の詰め合わせ問題を解く。
    N個のお菓子から長さK以下の連続する区間をM個選び、その和の合計を最大化する。
    """
    # 標準入力からすべてのデータを読み込み、スペースで分割する
    # N=3000程度なので、入力サイズは大きくなく、一括読み込みが効率的
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    it = iter(input_data)
    try:
        # N: お菓子の数, M: 作るセットの数, K: 1セットのお菓子の最大個数
        N = int(next(it))
        M = int(next(it))
        K = int(next(it))
        # 各お菓子のおいしさ
        A = list(map(int, it))
    except StopIteration:
        return
    
    # 累積和を計算する: P[i] は A[0]...A[i-1] の和
    # これにより、区間 [l, r] の和を P[r] - P[l-1] で O(1) で計算できる
    P = [0] * (N + 1)
    s = 0
    for i in range(N):
        s += A[i]
        P[i+1] = s
    
    # すべての可能な区間の和を格納するリスト
    all_sums = []
    # 速度向上のためにメソッドをローカル変数にキャッシュ
    all_sums_extend = all_sums.extend
    sub = operator.sub
    
    # 長さ 1 から K までのすべての連続する区間の和を計算する
    # 各長さについて、スライシングと map を組み合わせて効率的に計算
    for length in range(1, K + 1):
        # P[length:] は P[length], P[length+1], ..., P[N]
        # P は P[0], P[1], ..., P[N]
        # map(sub, P[length:], P) は P[length]-P[0], P[length+1]-P[1], ..., P[N]-P[N-length] を生成する
        # これは長さ 'length' のすべての区間の和に対応する
        all_sums_extend(map(sub, P[length:], P))
    
    # 区間の和を降順(大きい順)にソートする
    # Pythonの sort (Timsort) は非常に高速であり、4.5e6 要素程度なら1秒以内で処理可能
    all_sums.sort(reverse=True)
    
    # 最大の満足度を得るために、上位 M 個の区間の和を選択する
    # メモリ節約のため、M個より多い分は削除してから合計を計算する
    if len(all_sums) > M:
        del all_sums[M:]
    
    # 合計値を出力
    print(sum(all_sums))

if __name__ == '__main__':
    solve()

This editorial was generated by gemini-3-flash-thinking.

posted:
last update: