在长度不超过k的子数组中，找到和为s的子数组数量。

Question

我遇到了以下问题：

给定一个整数数组arr，一个正整数k和一个整数s，你的任务是找到长度不大于k且和等于s的非空连续子数组的数量。

对于arr = [1,2,4,-1,6,1]，k = 3和s = 6，输出应为solution(arr，k，s) = 3。

- 在长度为1的连续子数组中，有1个子数组与和等于s = 6，它是[6]， - 在长度为2的连续子数组中，有1个子数组与和等于s = 6，它是[2,4]， - 在长度为3的连续子数组中，有1个子数组与和等于s = 6，它是[-1,6,1]。

请注意，子数组[1,2,4，1]的和也为s，但它的长度大于k，因此不适用。

所以答案是3。

以下是我在Python中的尝试。我的想法是将每个前缀和的出现索引存储在字典中。例如，如果前缀和6在索引1和3处出现，则字典中的条目将是6：[1,3]。然后，在每一步中，我们可以检查是否已经遇到了目标和s（if curr - s in d:）以及子数组的范围。

这通过了10/15个测试用例，但超过了余下的隐藏用例的时间限制。如果有人有优化算法，我将不胜感激。

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

- codeedoc

---

这个循环是问题所在：for idx in d[curr-s]:，你可以通过在字典中存储的索引上进行二分查找来消除它。 - Photon

Answer 1

我们使用计数字典来跟踪前缀和的计数。我们还将使用大小为k的滑动窗口，根据前面的索引增加前缀和的计数，并根据后面的索引减少它。在我们减少之前，我们将根据累积和的计数添加到有效子数组的计数中，这些累积和比我们的后面累积和多k。

  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

- Dave