Official

E - 区間の合計 / Sum of Intervals Editorial by admin

Claude 4.5 Opus

Overview

This problem asks you to find the number of ways to select a contiguous subsequence from a sequence such that its sum is exactly \(K\). It can be solved efficiently by combining prefix sums with a hash map.

Analysis

Naive Approach and Its Issues

The simplest method is to try all contiguous subsequences: - Start point \(l\) from \(1\) to \(N\) - End point \(r\) from \(l\) to \(N\) - Calculate the sum of each interval

With this method, there are \(O(N^2)\) ways to choose intervals, and calculating each interval’s sum takes \(O(N)\), resulting in \(O(N^3)\) overall. This won’t work in time for \(N = 2 \times 10^5\).

Optimization with Prefix Sums

We introduce prefix sums \(S\): - \(S[0] = 0\) - \(S[i] = A[1] + A[2] + \cdots + A[i]\)

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

Key Insight

The condition for the sum of interval \([l, r]\) to equal \(K\) is: $\(S[r] - S[l-1] = K\)$

Rearranging this: $\(S[l-1] = S[r] - K\)$

In other words, for each \(r\), we need to quickly count how many prefix sums equal \(S[r] - K\).

Concrete Example

For sequence \(A = [1, 2, 3, -2, 2]\) with \(K = 3\): - \(S = [0, 1, 3, 6, 4, 6]\)

When \(r = 2\) (\(S[2] = 3\)), we look for \(S[l-1]\) where \(S[r] - K = 0\), and find \(S[0] = 0\) matches. This means the sum of interval \([1, 2]\) is \(3\).

Algorithm

  1. Prepare a hash map prefix_count and initialize with prefix_count[0] = 1 (corresponding to \(S[0] = 0\))
  2. Initialize cumulative sum current_sum to \(0\)
  3. Traverse array \(A\) from the beginning:
    • current_sum += A[i] (this corresponds to \(S[i+1]\))
    • Calculate target = current_sum - K
    • Add prefix_count[target] to the answer (the count of past prefix sums equal to target)
    • Increment prefix_count[current_sum] by \(1\)
  4. Output the answer
Example: A = [1, 2, -1, 2], K = 3

i=0: current_sum=1, target=-2, count+=0, prefix_count={0:1, 1:1}
i=1: current_sum=3, target=0,  count+=1, prefix_count={0:1, 1:1, 3:1}
i=2: current_sum=2, target=-1, count+=0, prefix_count={0:1, 1:1, 3:1, 2:1}
i=3: current_sum=4, target=1,  count+=1, prefix_count={0:1, 1:1, 3:1, 2:1, 4:1}

Answer: 2 (intervals [1,2] and [2,4])

Complexity

  • Time Complexity: \(O(N)\)
    • Traverse the array once, performing hash map operations (average \(O(1)\)) for each element
  • Space Complexity: \(O(N)\)
    • Store at most \(N+1\) prefix sums in the hash map

Implementation Notes

  • Initialization: Don’t forget to set prefix_count[0] = 1. This corresponds to the “select nothing” state (\(S[0] = 0\)) and is necessary to correctly count subsequences starting from the beginning
  • Order: Count the occurrences of target before adding the current prefix sum to prefix_count. This prevents counting itself as a duplicate
  • Handling Negative Numbers: Since elements can be negative, the two-pointer technique cannot be used. The hash map approach is appropriate

Source Code

from collections import defaultdict

def solve():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    
    # Using prefix sums
    # S[i] = A[0] + A[1] + ... + A[i-1] (S[0] = 0)
    # Sum of interval [l, r] = S[r+1] - S[l]
    # Condition for this to equal K: S[r+1] - S[l] = K
    # That is, S[l] = S[r+1] - K
    
    # Record how many times each prefix sum value has appeared
    prefix_count = defaultdict(int)
    prefix_count[0] = 1  # S[0] = 0
    
    current_sum = 0
    count = 0
    
    for i in range(N):
        current_sum += A[i]  # This corresponds to S[i+1]
        # How many prefix sums equal S[i+1] - K
        target = current_sum - K
        count += prefix_count[target]
        # Record the current prefix sum
        prefix_count[current_sum] += 1
    
    print(count)

solve()

This editorial was generated by claude4.5opus.

posted:
last update: