Leetcode 930 Binary Subarrays With Sum

发表于 |

给定一个只含0,1的数组，问数组中有多少个子串和恰好为S。

Example 1:
Input: A = [1,0,1,0,1], S = 2
Output: 4
Explanation:
The 4 subarrays are bolded below:
[1,0,1]
[1,0,1,0]
[0,1,0,1]
[1,0,1]

Note:

  1. A.length <= 30000
  2. 0 <= S <= A.length
  3. A[i] is either 0 or 1.

题意分析：
求数组中的连续子串，使其和恰好为S。
问这样的连续子串一共有多少个？

思路分析：
这和我之前做过的一道腾讯笔试题一模一样，链接如下，第一题(包含k个1的子串)：
腾讯笔试

关键就是在于记录数组中1的下标，
例如对于A = '001010011',S = 2，使用Index存储所有1的下标，那么有
index = [2,4,7,8]
假设我们选中前两个1(index1=2,index2=4)，这两个’1’一共能组成3*3=9种

3是怎么来的呢？是观察能向两边扩展多少个’0’。
第一个’1’左边有两个’0’，那么可以选择扩展(0个，1个，2个)，共三种情况
第二个’1’右边有两个’0’，同理，所以最后是3*3 = 9

中间0的个数是通过1的下标相减得到！
为了处理边界情况，我们给index前后各加一个边界(-1,n)

单独处理一下S = 0的情况，因为这种情况下选不了1

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# 60ms
class Solution(object):
    def numSubarraysWithSum(self, A, S):
        """
        :type A: List[int]
        :type S: int
        :rtype: int
        """
        if not A: return 0
        index = [-1]
        for i in range(len(A)):
            if A[i] == 1: index.append(i)
        index.append(len(A))

        res = 0
        if S == 0: # 单独处理S = 0的情况
            for i in range(1, len(index)):
                num = index[i] - index[i-1] - 1
                res += (num+1)*num // 2
            return res
        # 从index = 1开始遍历
        i,j = 1,1
        while i < len(index)-1:
            j = i + S
            if j >= len(index): break
            res += (index[i]-index[i-1]) * (index[j]-index[j-1])
            i += 1
        return res