公式

D - 連続練習日数 / Consecutive Practice Days 解説 by admin

Claude 4.5 Opus

Overview

This problem asks us to efficiently count the number of contiguous subarrays of length at least \(K\) whose sum is at least \(M\). We solve it by combining prefix sums with a Binary Indexed Tree (BIT).

Analysis

Naive Approach and Its Problems

If we try all pairs \((l, r)\), there are \(O(N^2)\) combinations. Furthermore, calculating the interval sum for each pair takes \(O(N)\), resulting in \(O(N^3)\) overall. This is far too slow for \(N \leq 2 \times 10^5\).

Speedup Using Prefix Sums

Interval sums can be calculated in \(O(1)\) using prefix sums. Define the prefix sum \(S\) as follows: - \(S[0] = 0\) - \(S[i] = A_1 + A_2 + \cdots + A_i\)

Then, the sum of interval \([l, r]\) can be obtained as \(S[r] - S[l-1]\).

Reformulating the Conditions

Let’s organize the conditions. Setting \(l' = l - 1\): - \(r - l + 1 \geq K\) is equivalent to \(r \geq l' + K\) - \(S[r] - S[l'] \geq M\) is equivalent to \(S[l'] \leq S[r] - M\)

In other words, for each \(r\), we need to count the number of \(l'\) satisfying “\(l' \leq r - K\) and \(S[l'] \leq S[r] - M\)”.

Idea for Further Speedup

We increase \(r\) from \(K\) to \(N\) sequentially, counting \(l'\) values that satisfy the conditions.

  • When \(r = K\), the only usable \(l'\) is \(0\)
  • When \(r = K+1\), usable \(l'\) values are \(0, 1\)

In other words, each time \(r\) increases by 1, we can add \(l' = r - K\) to the “candidate set”.

To efficiently count elements in the candidate set satisfying \(S[l'] \leq S[r] - M\), a BIT (Binary Indexed Tree) is effective.

Algorithm

  1. Calculate prefix sums: Compute \(S[0], S[1], \ldots, S[N]\)
  2. Coordinate compression: Since \(S\) values can reach approximately \(10^{14}\), compress coordinates to make them manageable for BIT
  3. Counting with BIT:
    • Iterate \(r\) from \(K\) to \(N\)
    • For each \(r\):
      • Add \(S[l']\) corresponding to \(l' = r - K\) to the BIT
      • Find how many values in the BIT are at most \(S[r] - M\) (= number of \(l'\) satisfying the condition) using binary search + BIT
    • Add to the answer

Concrete Example (\(N=5, K=2, M=10, A=[3,5,4,2,6]\)): - \(S = [0, 3, 8, 12, 14, 20]\) - \(r=2\): Add \(l'=0\). \(S[2]-M=8-10=-2\). Number of \(l'\) with \(S[l'] \leq -2\) is 0 - \(r=3\): Add \(l'=1\). \(S[3]-M=12-10=2\). \(S[l'] \leq 2\) holds for \(S[0]=0\), so 1 - …and so on

Complexity

  • Time complexity: \(O(N \log N)\)
    • Prefix sums: \(O(N)\)
    • Coordinate compression (sorting): \(O(N \log N)\)
    • BIT operations for each \(r\): \(O(\log N)\) × \(N\) times
  • Space complexity: \(O(N)\)
    • Prefix sum array, BIT, coordinate compression array

Implementation Notes

  • Coordinate compression: Convert \(S\) values to consecutive integers starting from \(1\) for use as BIT indices
  • Binary search: Use bisect_right to find “the maximum index at most \(S[r] - M\)”. Since bisect_right returns “the first position greater than the target”, we take the prefix sum up to that position
  • 1-indexed BIT: BIT is typically implemented as 1-indexed, so we use i + 1 during coordinate compression
  • Overflow prevention: Since \(M\) and \(A_i\) can be large, this isn’t an issue in Python, but use long long etc. in other languages

Source Code

import bisect

def solve():
    N, K, M = map(int, input().split())
    A = list(map(int, input().split()))
    
    # Calculate prefix sums
    # 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]
    
    # Condition: r - l + 1 >= K and S[r] - S[l-1] >= M
    # Consider l, r as 1-indexed
    # S[r] - S[l-1] >= M
    # r >= l + K - 1 (i.e., r - (l-1) >= K)
    
    # Let l' = l - 1 (0 <= l' <= N-K)
    # S[r] - S[l'] >= M
    # r >= l' + K
    # Range of r: l' + K <= r <= N
    
    # For each l', count the number of r where S[r] >= S[l'] + M
    # Range of r is [l' + K, N]
    
    # For efficiency, we look from the right side
    # Fix r and count l' where S[l'] <= S[r] - M and l' <= r - K
    
    # Can use BIT or segment tree,
    # here we use coordinate compression + BIT
    
    # Manage S[l'] and for each r, count the number where S[l'] <= S[r] - M
    # However, l' must be at most r - K
    
    # Move r from K to N
    # When r = K, only l' = 0 is possible
    # As r increases, add l' = r - K
    
    # Coordinate compression
    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
    
    # To find the count of S[l'] <= S[r] - M,
    # binary search for the maximum index <= S[r] - M
    
    count = 0
    
    for r in range(K, N + 1):
        # Add l' = r - K to BIT
        l_prime = r - K
        idx = compress[S[l_prime]]
        bit_add(idx)
        
        # Count S[l'] values <= S[r] - M
        target = S[r] - M
        # Find the maximum index in sorted_values that is <= target
        pos = bisect.bisect_right(sorted_values, target)
        if pos > 0:
            count += bit_sum(pos)
    
    print(count)

solve()

This editorial was generated by claude4.5opus.

投稿日時:
最終更新: