D - 連続練習日数 / Consecutive Practice Days Editorial by admin
Claude 4.5 Opus概要
長さ \(K\) 以上の連続部分列で、合計が \(M\) 以上となるものの個数を効率的に数える問題です。累積和と Binary Indexed Tree (BIT) を組み合わせて解きます。
考察
素朴なアプローチとその問題点
すべての \((l, r)\) の組を試すと、\(O(N^2)\) 通りの組み合わせがあります。さらに各組について区間和を計算すると \(O(N)\) かかり、全体で \(O(N^3)\) となってしまいます。\(N \leq 2 \times 10^5\) では到底間に合いません。
累積和による高速化
区間和は累積和を使えば \(O(1)\) で計算できます。累積和 \(S\) を以下のように定義します: - \(S[0] = 0\) - \(S[i] = A_1 + A_2 + \cdots + A_i\)
すると、区間 \([l, r]\) の和は \(S[r] - S[l-1]\) で求められます。
条件の言い換え
条件を整理しましょう。\(l' = l - 1\) と置くと: - \(r - l + 1 \geq K\) は \(r \geq l' + K\) と同値 - \(S[r] - S[l'] \geq M\) は \(S[l'] \leq S[r] - M\) と同値
つまり、各 \(r\) に対して「\(l' \leq r - K\) かつ \(S[l'] \leq S[r] - M\)」を満たす \(l'\) の個数を数えればよいです。
さらなる高速化の発想
\(r\) を \(K\) から \(N\) まで順に増やしながら、条件を満たす \(l'\) を数えます。
- \(r = K\) のとき、使える \(l'\) は \(0\) のみ
- \(r = K+1\) のとき、使える \(l'\) は \(0, 1\)
- …
つまり、\(r\) が1増えるごとに、\(l' = r - K\) を「候補集合」に追加できます。
候補集合の中から \(S[l'] \leq S[r] - M\) を満たすものを高速に数えるには、BIT(Binary Indexed Tree) が有効です。
アルゴリズム
- 累積和の計算: \(S[0], S[1], \ldots, S[N]\) を計算
- 座標圧縮: \(S\) の値は最大 \(10^{14}\) 程度になるため、BIT で扱えるよう座標圧縮
- BIT による計数:
- \(r\) を \(K\) から \(N\) まで順に走査
- 各 \(r\) について:
- \(l' = r - K\) に対応する \(S[l']\) を BIT に追加
- \(S[r] - M\) 以下の値が BIT に何個あるか(= 条件を満たす \(l'\) の個数)を二分探索 + BIT で求める
- 答えに加算
具体例 (\(N=5, K=2, M=10, A=[3,5,4,2,6]\)): - \(S = [0, 3, 8, 12, 14, 20]\) - \(r=2\): \(l'=0\) を追加。\(S[2]-M=8-10=-2\)。\(S[l'] \leq -2\) となる \(l'\) は 0個 - \(r=3\): \(l'=1\) を追加。\(S[3]-M=12-10=2\)。\(S[l'] \leq 2\) となるのは \(S[0]=0\) の1個 - …と続く
計算量
- 時間計算量: \(O(N \log N)\)
- 累積和: \(O(N)\)
- 座標圧縮(ソート): \(O(N \log N)\)
- 各 \(r\) に対する BIT 操作: \(O(\log N)\) × \(N\) 回
- 空間計算量: \(O(N)\)
- 累積和配列、BIT、座標圧縮用配列
実装のポイント
座標圧縮: BIT の添字として使うため、\(S\) の値を \(1\) から始まる連番に変換
二分探索:
bisect_rightで「\(S[r] - M\) 以下の最大インデックス」を求める。bisect_rightは「より大きい最初の位置」を返すので、その位置までの累積和を取るBIT の 1-indexed: BIT は通常 1-indexed で実装するため、座標圧縮時に
i + 1としているオーバーフロー対策: \(M\) や \(A_i\) が大きいため、Python では問題ないが他言語では
long long等を使用ソースコード
import bisect
def solve():
N, K, M = map(int, input().split())
A = list(map(int, input().split()))
# 累積和を計算
# S[i] = A[0] + A[1] + ... + A[i-1] (S[0] = 0)
S = [0] * (N + 1)
for i in range(N):
S[i + 1] = S[i] + A[i]
# 条件: r - l + 1 >= K かつ S[r] - S[l-1] >= M
# l, r は 1-indexed で考える
# S[r] - S[l-1] >= M
# r >= l + K - 1 (つまり r - (l-1) >= K)
# l' = l - 1 とすると (0 <= l' <= N-K)
# S[r] - S[l'] >= M
# r >= l' + K
# r の範囲: l' + K <= r <= N
# 各 l' に対して、S[r] >= S[l'] + M となる r の数を数える
# r の範囲は [l' + K, N]
# 効率化のため、右端から見ていく
# r を固定して、S[l'] <= S[r] - M かつ l' <= r - K となる l' の数を数える
# BIT や セグメント木を使う方法もあるが、
# ここでは座標圧縮 + BIT を使用
# S[l'] を管理し、各 r に対して S[l'] <= S[r] - M となる個数を数える
# ただし l' は r - K 以下でなければならない
# r を K から N まで動かす
# r = K のとき、l' は 0 のみ可能
# r を増やすごとに、l' = r - K を追加していく
# 座標圧縮
all_values = set(S)
sorted_values = sorted(all_values)
compress = {v: i + 1 for i, v in enumerate(sorted_values)} # 1-indexed for BIT
size = len(sorted_values)
# BIT (Binary Indexed Tree)
bit = [0] * (size + 2)
def bit_add(i, x=1):
while i <= size:
bit[i] += x
i += i & (-i)
def bit_sum(i):
s = 0
while i > 0:
s += bit[i]
i -= i & (-i)
return s
# S[r] - M 以下の S[l'] の個数を求めるには、
# S[r] - M 以下の最大のインデックスを二分探索で求める
count = 0
for r in range(K, N + 1):
# l' = r - K を BIT に追加
l_prime = r - K
idx = compress[S[l_prime]]
bit_add(idx)
# S[r] - M 以下の S[l'] の個数を数える
target = S[r] - M
# sorted_values の中で target 以下の最大のインデックスを探す
pos = bisect.bisect_right(sorted_values, target)
if pos > 0:
count += bit_sum(pos)
print(count)
solve()
この解説は claude4.5opus によって生成されました。
posted:
last update: