Find number of subarrays with length

Question

I encountered the following question:

Given an array of integers arr, a positive integer k, and an integer s, your task is to find the number of non-empty contiguous subarrays with length not greater than k and with a sum equal to s.

For arr = [1, 2, 4, -1, 6, 1], k = 3, and s = 6, the output should be solution(arr, k, s) = 3.

There is 1 subarray among the contiguous subarrays of length 1 with sum equal to s = 6, and it is [6],
There is 1 subarray among the contiguous subarrays of length 2 with sum equal to s = 6, and it is [2, 4],
There is 1 subarray among the contiguous subarrays of length 3 with sum equal to s = 6, and it is [-1, 6, 1].

Note that the subarray [1, 2, 4, -1] would also sum to s, but its length is greater than k, so it's not applicable.

So the answer is 3.

Below is my attempt in Python. My idea is to store the indices of each prefix sum's occurrences in a dictionary. For example, if a prefix sum 6 occurs at indices 1 and 3, the entry in the dictionary would be 6:[1, 3]. Then at each step we could check if the target sum s has encountered (if curr - s in d:) and the range of the subarray.

This passed 10/15 of the test cases but went over the time limit for the remaining hidden cases. I'd appreciate if anyone has an optimized algorithm.

import collections
def solution(arr, k, s):
    res = 0
    curr = 0
    d = collections.defaultdict(list)
    d[0].append(-1)
    for i, num in enumerate(arr):
        curr += num
        if curr - s in d:
            for idx in d[curr-s]:
                if i - idx <= k:
                    res += 1
        d[curr].append(i)
    return res

user_2041077 · Answer

We use a counting dictionary to keep track of the count of prefix sums. We'll also use a sliding window of size k, incrementing the count of prefix sums based on the front index, and decrementing it based on the rear index. Right before we decrement, we'll add to our count of valid subarrays based on the count of cumulative sums that, are k more than our rear cumulative sum.

  def solution(arr, k, s)
    num_subarrays = 0
    front_cumulative_sum = 0
    rear_cumulative_sum = 0
    cumulative_sum_to_count = Hash.new { |h, cumsum| h[cumsum] = 0 }
    0.upto(arr.size - 1 + k) do |front_index|
      front_cumulative_sum += arr[front_index] if front_index <= arr.size - 1
      cumulative_sum_to_count[front_cumulative_sum] += 1 if front_index <= arr.size - 1
      if front_index >= k
        rear_index = front_index - k
        rear_cumulative_sum += arr[rear_index]
        num_subarrays += cumulative_sum_to_count[s + rear_cumulative_sum]
        cumulative_sum_to_count[rear_cumulative_sum] -= 1
      end
    end
    return num_subarrays
  end


> arr = [1,2,4,-1,6,1]
> k=3
> s=6
> solution(arr, k, s)
=> 3

Find number of subarrays with length <= k and with sum == s

1 Answers