LeetCode-in-All

560. Subarray Sum Equals K

Medium

Given an array of integers nums and an integer k, return the total number of subarrays whose sum equals to k.

A subarray is a contiguous non-empty sequence of elements within an array.

Example 1:

Input: nums = [1,1,1], k = 2

Output: 2

Example 2:

Input: nums = [1,2,3], k = 3

Output: 2

Constraints:

Solution

(define/contract (subarray-sum nums k)
  (-> (listof exact-integer?) exact-integer? exact-integer?)
  (let ([temp-sum 0]
        [result 0]
        [sum-count (make-hash)])
    
    ; Initialize the hash with 0 sum occurring once
    (hash-set! sum-count 0 1)
    
    ; Process each number
    (for ([num nums])
      ; Update running sum
      (set! temp-sum (+ temp-sum num))
      
      ; Check if we have seen temp-sum - k before
      (when (hash-has-key? sum-count (- temp-sum k))
        (set! result (+ result (hash-ref sum-count (- temp-sum k)))))
      
      ; Update the count for current sum
      (hash-set! sum-count 
                 temp-sum 
                 (add1 (hash-ref sum-count temp-sum 0))))
    
    result))