Official

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) が有効です。

アルゴリズム

  1. 累積和の計算: \(S[0], S[1], \ldots, S[N]\) を計算
  2. 座標圧縮: \(S\) の値は最大 \(10^{14}\) 程度になるため、BIT で扱えるよう座標圧縮
  3. 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: