Skip to content

How to Find All Subarrays With Given Sum in Python

Updated: at 04:23 AM

Given an array of integers and a target sum k, the subarray sum problem involves finding the total count of contiguous subarrays having a sum equal to the given target sum. This is a common interview question asked during technical interviews for software engineering or data science roles.

Mastering this concept requires knowledge of two fundamental programming techniques - sliding windows and prefix sums. In this comprehensive Python programming guide, we will provide:

Table of Contents

Open Table of Contents

Formal Problem Statement

Given an unsorted array nums of n integers and a target sum k, return the total count of continuous subarrays that have a sum equal to k.

Example 1:

Input: nums = [1,1,1], k = 2
Output: 2
Explanation: The subarrays [1,1] and [1,1] have a sum of 2.

Example 2:

Input: nums = [1,2,3], k = 3
Output: 2
Explanation: The subarrays [1,2] and [3] have a sum of 3.

Brute Force Approach

The most straightforward technique is using a brute force approach - check the sum of all possible subarrays to see if any of them equal k.

Here are the steps:

  1. Use two for loops to generate all subarrays:
    • Outer loop picks the starting index i, from 0 to n-1
    • Inner loop picks ending index j, from i+1 to n
  2. For each subarray [i...j], calculate the sum and compare it with target k
  3. If the sum equals k, increment the count
  4. Return the final count

Python Implementation:

def subarray_sum_bruteforce(nums, k):

  count = 0

  for i in range(len(nums)):
    for j in range(i+1, len(nums)+1):

      # Generate subarray
      subarray = nums[i:j]

      # Calculate sum of subarray
      sub_sum = sum(subarray)

      # Compare sum with target
      if sub_sum == k:
        count += 1

  return count

This brute force technique works by generating all subarrays one by one and checking if any of them have the desired sum k.

Time Complexity - O(N^2) since we have nested loops

Space Complexity - O(1) since only constant extra space is used

The brute force method is simple to implement but inefficient for large inputs. We need a more optimal approach.

Optimized Solution - Sliding Window

We can optimize this using the sliding window technique. The steps are:

  1. Initialize a variable window_sum to keep track of current window’s sum

  2. Initialize left and right pointers l and r at index 0

  3. Loop through the array:

    • Add current element to window_sum
    • If window_sum is greater than k, subtract elements from the left to shrink window
    • If window_sum equals k, increment count and shrink the window to exclude the leftmost element
    • Increment the right pointer to continue sliding window
  4. Return final count

This ensures that the window size is minimized, reducing redundant additions and comparisons.

Python Implementation:

def subarray_sum_slidingwindow(nums, k):

  count = 0
  window_sum = 0

  l = 0

  for r in range(len(nums)):

    # Add current element
    window_sum += nums[r]

    # Shrink window if outside k
    while window_sum > k and l < r:
      window_sum -= nums[l]
      l += 1

    # Found subarray with k sum
    if window_sum == k:
      count += 1

      # Shrink left bound
      window_sum -= nums[l]
      l += 1

  return count

The sliding window approach optimizes this to O(N) time complexity since we iterate through the array just once.

Space complexity is O(1) as well. This solution is much more efficient than brute force.

Prefix Sum Approach

The prefix sum technique is another method we can apply here. The steps are:

  1. Generate prefix sum array prefix where prefix[i] is sum of elements nums[0...i]

  2. Use a hashmap map to store seen prefix sums and their counts

  3. Loop through prefix array:

    • Calculate current required prefix sum as prefix[i] - k
    • If this required sum already exists in map, there is a subarray ending at i with sum k
    • Increment count by no. of times the required prefix sum has been seen
    • Update map for current prefix[i]
  4. Return final count

This works because if a subarray [i...j] has sum k, then prefix[j] - prefix[i-1] = k. So required prefix sums can be calculated.

Python Implementation:

def subarray_sum_prefixsum(nums, k):

  count = 0
  prefix = [0] * (len(nums) + 1)
  map_sum = {0:1}

  # Generate prefix sum
  for i in range(1, len(prefix)):
    prefix[i] = prefix[i-1] + nums[i-1]

  # Check map for required sum
  for i in range(len(prefix)):

    required_sum = prefix[i] - k
    if required_sum in map_sum:
      count += map_sum[required_sum]

    map_sum[prefix[i]] = map_sum.get(prefix[i], 0) + 1

  return count

This prefix sum method also has O(N) time complexity since we traverse the nums array only once.

Space complexity is O(N) due to the prefix sum array. This provides another optimal solution.

Analysis

Let’s summarize the time and space complexities of the approaches discussed:

AlgorithmTime ComplexitySpace Complexity
Brute ForceO(N^2)O(1)
Sliding WindowO(N)O(1)
Prefix SumO(N)O(N)

The sliding window and prefix sum techniques are clearly more efficient than brute force in terms of time complexity.

Space complexity is constant for sliding window since we only use a few variables. Prefix sum uses O(N) extra space to store the prefix array.

So in conclusion, sliding window provides the optimal approach for solving this problem in linear time with minimal space usage.

Practical Examples

Here are some examples of how this subarray sum algorithm can be useful:

Developers can easily adapt this algorithm to various applications involving array or sequence analysis.

Summary

In this comprehensive guide, we looked at various techniques to find all subarrays with a given sum k in Python:

By mastering these algorithms, you can confidently tackle subarray problems in interviews and programming projects. The concepts of sliding windows and prefix sums are also broadly applicable.

Hopefully this step-by-step article provided clear explanations of the logic and Python code needed to find subarrays with a given sum. Let us know if you have any other questions!