LeetCode 560. Subarray Sum Equals K

Post by ailswan Aug. 31, 2024

560. Subarray Sum Equals K

Problem Statement

link: LeetCode.cn LeetCode 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:

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

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

Constraints: 1 <= nums.length <= 2 * 10^4 -1000 <= nums[i] <= 1000 -10^7 <= k <= 10^7

Solution Approach

The solution utilizes a prefix sum approach with a hash table to efficiently count the number of subarrays with a sum equal to k by tracking cumulative sums and their frequencies.

Algorithm

  1. Prefix Sum Calculation: The algorithm uses a running prefix sum to track cumulative sums of elements.
  2. Hash Table Utilization: It employs a hash table to efficiently store and query prefix sums to find subarrays that sum to k.
  3. Initial Setup: The hash table is initialized with {0: 1} to handle subarrays starting from the beginning.

Complexity

time complexity: O(n) space complexity: O(n)

Implement

    class Solution:
   def subarraySum(self, nums: List[int], k: int) -> int:
       pre_sums = defaultdict(int)
       res = 0
       pre_s = 0
       pre_sums[0] = 1
       for i, n in enumerate(nums):
           pre_s += n
           if pre_s - k in pre_sums:
               res += pre_sums[pre_s - k]
           pre_sums[pre_s] += 1
       return res