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
- Prepare prefix sums: Compute the prefix sum array \(P\) of sequence \(A\).
- Enumerate all interval sums: For all contiguous intervals of length \(1\) to \(K\), compute their sums and store them in a list.
- Sort: Sort the stored interval sums in descending order (largest first).
- 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
forloop withappendone by one, leveragingmap, slicing, andextendcan 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]usingoperator.subandmap, 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: