Taro Logo

Binary Subarrays With Sum

Medium
2 views
2 months ago

Given a binary array nums and an integer goal, return the number of non-empty subarrays with a sum goal.

A subarray is a contiguous part of the array.

Example 1:

Input: nums = [1,0,1,0,1], goal = 2
Output: 4
Explanation: The 4 subarrays are bolded and underlined below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

Example 2:

Input: nums = [0,0,0,0,0], goal = 0
Output: 15

Consider the following:

  1. How would you solve this problem using a brute force approach? What is the time and space complexity?
  2. How would you optimize the solution? What is the time and space complexity of the optimal solution?
  3. Can you walk through the optimal solution with the example nums = [1,0,1,0,1] and goal = 2?
  4. What edge cases should you consider and how does your solution handle them? (e.g. empty array, all zeros, goal is zero, goal is larger than the sum of all elements)
Sample Answer
## Problem: Number of Subarrays with Given Sum

Given a binary array `nums` and an integer `goal`, return the number of non-empty subarrays with a sum `goal`.

A subarray is a contiguous part of the array.

**Example 1:**

Input: nums = [1,0,1,0,1], goal = 2 Output: 4 Explanation: The 4 subarrays are bolded and underlined below: [1,0,1,0,1] [1,0,1,0,1] [1,0,1,0,1] [1,0,1,0,1]


**Example 2:**

Input: nums = [0,0,0,0,0], goal = 0 Output: 15


## Solution

### 1. Brute Force Approach

The simplest approach is to iterate through all possible subarrays and check if their sum equals the given `goal`.  This involves two nested loops to define the start and end indices of the subarray, and an inner loop (or a sum calculation within the outer loops) to calculate the sum of the subarray.

```python
def num_subarrays_with_sum_brute_force(nums, goal):
    count = 0
    n = len(nums)
    for i in range(n):
        for j in range(i, n):
            subarray_sum = sum(nums[i:j+1])
            if subarray_sum == goal:
                count += 1
    return count

2. Optimal Approach: Prefix Sum and Hashmap

We can optimize this by using a prefix sum and a hashmap. The prefix sum at index i is the sum of all elements from index 0 to i. We store the prefix sums in a hashmap, where the key is the prefix sum and the value is the number of times that prefix sum has occurred.

To find the number of subarrays with sum goal, we iterate through the array, calculate the prefix sum at each index, and check if prefix_sum - goal exists in the hashmap. If it does, it means there is a subarray ending at the current index with a sum equal to goal. We add the number of times prefix_sum - goal has occurred to the count.

def num_subarrays_with_sum(nums, goal):
    prefix_sum = 0
    count = 0
    prefix_sum_count = {0: 1}

    for num in nums:
        prefix_sum += num
        if prefix_sum - goal in prefix_sum_count:
            count += prefix_sum_count[prefix_sum - goal]
        
        if prefix_sum in prefix_sum_count:
            prefix_sum_count[prefix_sum] += 1
        else:
            prefix_sum_count[prefix_sum] = 1

    return count

Example Walkthrough

Let's walk through nums = [1,0,1,0,1] and goal = 2.

  1. Initialize prefix_sum = 0, count = 0, prefix_sum_count = {0: 1}.
  2. First element 1: prefix_sum = 1. prefix_sum - goal = 1 - 2 = -1. -1 is not in prefix_sum_count. Update prefix_sum_count = {0: 1, 1: 1}.
  3. Second element 0: prefix_sum = 1. prefix_sum - goal = 1 - 2 = -1. -1 is not in prefix_sum_count. Update prefix_sum_count = {0: 1, 1: 2}.
  4. Third element 1: prefix_sum = 2. prefix_sum - goal = 2 - 2 = 0. 0 is in prefix_sum_count. count += prefix_sum_count[0] = 1. Update prefix_sum_count = {0: 1, 1: 2, 2: 1}.
  5. Fourth element 0: prefix_sum = 2. prefix_sum - goal = 2 - 2 = 0. 0 is in prefix_sum_count. count += prefix_sum_count[0] = 1. Update prefix_sum_count = {0: 1, 1: 2, 2: 2}.
  6. Fifth element 1: prefix_sum = 3. prefix_sum - goal = 3 - 2 = 1. 1 is in prefix_sum_count. count += prefix_sum_count[1] = 2. Update prefix_sum_count = {0: 1, 1: 2, 2: 2, 3: 1}.
  7. Return count = 4.

Big(O) Runtime Analysis

  • Brute Force: O(n^2) due to the nested loops to generate all subarrays, and potentially O(n^3) if you recalculate the sum of each subarray within the loops. If you keep a running sum while generating subarrays, it can be reduced to O(n^2).
  • Optimal (Prefix Sum with Hashmap): O(n), where n is the length of the input array nums. This is because we iterate through the array only once.

Big(O) Space Usage Analysis

  • Brute Force: O(1) as we are not using any extra space.
  • Optimal (Prefix Sum with Hashmap): O(n) in the worst case, where n is the length of the input array nums. This is because the prefix_sum_count hashmap can potentially store all the prefix sums if all the elements are unique.

Edge Cases

  1. Empty Array: If the input array is empty, the number of subarrays with sum goal will always be 0. The above solutions will handle this case correctly.
  2. All Zeros: If the input array contains all zeros and the goal is 0, then every subarray will have a sum of 0. The number of such subarrays is n * (n + 1) / 2, where n is the length of the array. The optimal solution will handle this correctly.
  3. Goal is Zero: If the goal is 0, we are essentially looking for subarrays with all zeros. The optimal solution using the prefix sum and hashmap handles this case correctly by counting the occurrences of each prefix sum.
  4. Goal is Greater than Sum of All Elements: If the goal is greater than the sum of all elements in the array, then no subarray can have the required sum. The optimal solution will handle this correctly as the count will remain zero.