DSA Crash Course
Templates

Template: 13. PREFIX SUM

Mark when done:

13. PREFIX SUM

# =============================================================================
def prefix_sum(arr):
    prefix = [0] * (len(arr) + 1)
    for i in range(len(arr)):
        prefix[i + 1] = prefix[i] + arr[i]
    # range sum [i..j] = prefix[j+1] - prefix[i]
    return prefix

def subarray_sum_equals_k(arr, k):
    """Count subarrays with sum = k using prefix sum + hashmap."""
    count = 0
    prefix = 0
    seen = {0: 1}  # prefix_sum -> count
    for num in arr:
        prefix += num
        if prefix - k in seen:
            count += seen[prefix - k]
        seen[prefix] = seen.get(prefix, 0) + 1
    return count


# =============================================================================