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

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

Answer 1

Make your dict as described, but also delete anything k from the front index. On deletion, add the count of good runs made with the index being deleted as the left end.

I'm on my phone. Conment if this is unclear and I can add details tomorrow evening.