In an array A of 0s and 1s, how many non-empty subarrays have sum S?
 
Example 1:
Input: A = [1,0,1,0,1], S = 2
Output: 4
Explanation: 
The 4 subarrays are bolded below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

 
Note:

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

Source: leetcode | Difficulty: MEDIUM
Tags: Data structures

Amortized analysis

Range queries

Starter code: class Solution:
    def numSubarraysWithSum(self, A: List[int], S: int) -> int:
        

Solutions:
class Solution:

	def numSubarraysWithSum(self, pl, S):
		ans = 0
		if S == 0:
			c = 0
			for i in range(len(pl)):
				if pl[i] == 0:
					c += 1
				else:
					c = 0
				ans += c
			return ans
		l = [-1]
		for i in range(len(pl)):
			if pl[i] == 1:
				l.append(i)
		l.append(len(pl))
		ans = 0
		for i in range(1, len(l) - S):
			ans += (l[i] - l[i - 1]) * (l[i + S] - l[i + S - 1])
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		size = len(A)
		(l, r, total, cnt) = [0] * 4
		if S == 0:
			l = 1
			for x in A:
				if x == 0:
					cnt += l
					l += 1
				else:
					l = 1
			return cnt
		while r < size and total < S:
			total += A[r]
			r += 1
		if r == size and total < S:
			return 0
		while r < size:
			(cnt1, cnt2) = (1, 1)
			while l < r and A[l] == 0:
				l += 1
				cnt1 += 1
			while r < size and A[r] == 0:
				r += 1
				cnt2 += 1
			cnt += cnt1 * cnt2
			l += 1
			r += 1
		if A[-1] == 1:
			cnt += 1
			while l < size and A[l] == 0:
				cnt += 1
				l += 1
		return cnt


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		P = [0]
		for x in A:
			P.append(P[-1] + x)
		count = collections.Counter()
		ans = 0
		for x in P:
			ans += count[x]
			count[x + S] += 1
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		n = len(A)
		suma = [0] * (n + 1)
		for i in range(1, n + 1):
			suma[i] = suma[i - 1] + A[i - 1]
		left = 0
		res = 0
		m = collections.defaultdict(lambda : 0)
		m[0] = 1
		for i in range(n):
			if suma[i + 1] - S in m:
				res += m[suma[i + 1] - S]
			m[suma[i + 1]] += 1
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:

		def atMostK(A, S):
			if S < 0:
				return 0
			i = res = 0
			for j in range(len(A)):
				S -= A[j]
				while S < 0:
					S += A[i]
					i += 1
				res += j - i + 1
			return res
		return atMostK(A, S) - atMostK(A, S - 1)


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		container = [0]
		for num in A:
			container.append(container[-1] + num)
		count = Counter()
		ans = 0
		for num in container:
			ans += count[num]
			count[num + S] += 1
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:

		def findZerosBefore(index, array):
			zerosBefore = 0
			index -= 1
			while index >= 0 and array[index] != 1:
				zerosBefore += 1
				index -= 1
			return zerosBefore

		def findZerosAfter(index, array):
			zerosAfter = 0
			index += 1
			while index < len(array) and array[index] != 1:
				zerosAfter += 1
				index += 1
			return zerosAfter

		def zeroString(index, array):
			zeros = 0
			while index < len(array) and array[index] != 1:
				zeros += 1
				index += 1
			return zeros

		def zeroCombs(n):
			combs = 0
			for i in range(n, 0, -1):
				combs = combs + n - i + 1
			return combs

		def firstZero(index, array):
			if index == 0:
				return True
			if array[index - 1] == 1:
				return True
			else:
				return False
		total = 0
		if S > 0:
			one_array = [i for (i, one) in enumerate(A) if one == 1]
			print(one_array)
			start_index = []
			end_index = []
			for n in range(len(one_array)):
				if n + S - 1 < len(one_array):
					start_index.append(one_array[n])
					end_index.append(one_array[n + (S - 1)])
			print(start_index)
			print(end_index)
			for i in range(len(start_index)):
				Bef = findZerosBefore(start_index[i], A)
				print(Bef)
				Aft = findZerosAfter(end_index[i], A)
				total = total + Bef + Aft + Bef * Aft + 1
		if S == 0:
			zero_array = [i for (i, zero) in enumerate(A) if zero == 0]
			for n in range(len(zero_array)):
				if firstZero(zero_array[n], A):
					sticky = zeroString(zero_array[n], A)
					total = int(total + (sticky ** 2 + sticky) / 2)
		return total


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		(i, s, res, count) = (0, 0, 0, 1)
		for j in range(len(A)):
			s += A[j]
			while s > S:
				s -= A[i]
				i += 1
				count = 1
			while i < j and A[i] == 0:
				i += 1
				count += 1
			if s == S and i <= j:
				res += count
		return res


from collections import Counter

class Solution:

	def numSubarraysWithSum(self, A, k):
		count = Counter({0: 1})
		ans = psum = 0
		for v in A:
			psum += v
			ans += count[psum - k]
			count[psum] += 1
		return ans


class Solution:

	def f(self, A):
		ret = 0
		cnt = 0
		for a in A:
			if not cnt:
				if not a:
					cnt += 1
			elif not a:
				cnt += 1
			else:
				ret += cnt * (cnt + 1) // 2
				cnt = 0
		if cnt:
			ret += cnt * (cnt + 1) // 2
		return ret
		pass

	def numSubarraysWithSum(self, A: List[int], k: int) -> int:
		ret = 0
		if not k:
			return self.f(A)
		mp = {0: 1}
		cnt = 0
		for a in A:
			cnt += a
			if cnt not in mp:
				mp[cnt] = 0
			mp[cnt] += 1
			if cnt - k in mp:
				ret += mp[cnt - k]
		return ret


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		h = {0: 1}
		count = summ = 0
		for i in range(0, len(A)):
			summ += A[i]
			if summ - S in h:
				count += h[summ - S]
			if summ not in h:
				h[summ] = 0
			h[summ] += 1
		return count


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		dp = []
		oneIndices = [-1]
		ans = 0
		for (i, num) in enumerate(A):
			if num == 1:
				oneIndices.append(i)
			if len(oneIndices) >= S + 1:
				if S == 0:
					ans += i - oneIndices[-S - 1]
				else:
					ans += oneIndices[-S] - oneIndices[-S - 1]
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		return self.at_most(A, S) - self.at_most(A, S - 1)

	def at_most(self, A, S):
		if S < 0:
			return 0
		start = 0
		res = 0
		for end in range(len(A)):
			S -= A[end]
			while start < len(A) and S < 0:
				S += A[start]
				start += 1
			res += end - start + 1
		return res


class Solution:

	def numSubarraysWithSum(self, nums: List[int], target: int) -> int:
		return self._at_most(nums, target) - self._at_most(nums, target - 1)

	def _at_most(self, nums, target):
		cnt = 0
		sums = 0
		i = 0
		for j in range(len(nums)):
			sums += nums[j]
			while i <= j and sums > target:
				sums -= nums[i]
				i += 1
			if sums <= target:
				cnt += j - i + 1
		return cnt


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		csum = [0]
		for a in A:
			csum.append(csum[-1] + a)
		counter = collections.Counter()
		ans = 0
		print(csum)
		for p in csum:
			ans += counter[p - S]
			counter[p] += 1
		return ans


import collections

class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		cache = collections.Counter({0: 1})
		currentSum = 0
		count = 0
		for num in A:
			currentSum += num
			count += cache[currentSum - S]
			cache[currentSum] += 1
		return count


class Solution:

	def numSubarraysWithSum(self, A, S):
		c = collections.Counter({0: 1})
		psum = res = 0
		for i in A:
			psum += i
			res += c[psum - S]
			c[psum] += 1
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		n = len(A)
		prefix = [0 for _ in range(n + 1)]
		for i in range(n):
			prefix[i + 1] = prefix[i] + A[i]

		def at_most_k(k: int) -> int:
			(begin, end) = (0, 1)
			cnt = 0
			while end < n + 1:
				while begin < end and prefix[end] - prefix[begin] > k:
					begin += 1
				cnt += end - begin
				end += 1
			return cnt
		return at_most_k(S) - at_most_k(S - 1)


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		return self.atMost(A, S) - self.atMost(A, S - 1)

	def atMost(self, A, S):
		res = 0
		i = 0
		pre_sum = 0
		for (j, val) in enumerate(A):
			pre_sum += val
			while pre_sum > S and i <= j:
				print(pre_sum, i)
				pre_sum -= A[i]
				i += 1
			res += j - i + 1
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:

		def atMost(S):
			(ans, l, tmp) = (0, 0, 0)
			for r in range(len(A)):
				tmp = tmp + A[r]
				while tmp > S and l <= r:
					tmp = tmp - A[l]
					l = l + 1
				ans += r - l + 1
			return ans
		return atMost(S) - atMost(S - 1)


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		left = res = cur_sum = count = 0
		for right in range(len(A)):
			cur_sum += A[right]
			if A[right] == 1:
				count = 0
			while left <= right and cur_sum >= S:
				if cur_sum == S:
					count += 1
				cur_sum -= A[left]
				left += 1
			res += count
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		cur = 0
		res = 0
		dp = collections.defaultdict(int)
		dp[0] = 1
		for i in range(len(A)):
			cur += A[i]
			diff = cur - S
			if diff in dp:
				res += dp[diff]
			dp[cur] += 1
		return res


class Solution:

	def numSubarraysWithSum(self, A, S):
		return self.numSubarraysAtMostSum(A, S) - self.numSubarraysAtMostSum(A, S - 1)

	def numSubarraysAtMostSum(self, A, S):
		if S < 0:
			return 0
		if not A:
			return 0
		n = len(A)
		right = 0
		now_sum = 0
		res = 0
		for left in range(n):
			while right <= n - 1 and (now_sum < S or A[right] == 0):
				now_sum += A[right]
				right += 1
			res += right - left
			now_sum -= A[left]
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		hashmap = {0: 1}
		preSum = 0
		ans = 0
		for i in range(len(A)):
			preSum += A[i]
			if preSum - S in hashmap:
				ans += hashmap[preSum - S]
			hashmap[preSum] = 1 if preSum not in hashmap else hashmap[preSum] + 1
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		if not A or S < 0:
			return 0
		map = Counter()
		map[0] = 1
		sum = 0
		res = 0
		for num in A:
			sum += num
			if sum - S in map:
				res += map[sum - S]
			map[sum] += 1
		return res


class Solution:

	def numSubarraysWithSum(self, nums: List[int], k: int) -> int:
		if not nums:
			return 0
		counter = collections.defaultdict(int)
		counter[0] = 1
		cumsum = 0
		ans = 0
		for i in range(len(nums)):
			cumsum += nums[i]
			if cumsum - k in counter:
				ans += counter[cumsum - k]
			counter[cumsum] += 1
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		size = len(A)
		prefix = {0: 1}
		cnt = 0
		ans = 0
		for (i, num) in enumerate(A):
			if num == 1:
				cnt += 1
			if cnt - S in prefix:
				ans += prefix[cnt - S]
			prefix[cnt] = prefix.get(cnt, 0) + 1
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		A = [0] + list(itertools.accumulate(A))
		hm = Counter()
		res = 0
		for i in A:
			res += hm[i]
			hm[i + S] += 1
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		count = {0: 1}
		prefix = 0
		result = 0
		for a in A:
			prefix = prefix + a
			if prefix - S in count:
				result += count[prefix - S]
			if prefix in count:
				count[prefix] += 1
			else:
				count[prefix] = 1
		return result


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		if S == 0 and sum(A) == 0:
			return len(A) * (len(A) + 1) // 2

		def atmost(s):
			if s < 0:
				return 0
			sumi = s
			l = 0
			r = 0
			count = 0
			while r < len(A):
				sumi -= A[r]
				while sumi < 0:
					sumi += A[l]
					l += 1
				count += r - l + 1
				r += 1
			return count
		return atmost(S) - atmost(S - 1)


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		d = {0: 1}
		ss = 0
		for i in A:
			ss += i
			d[ss] = d.get(ss, 0) + 1
		res = 0
		print(d)
		if S > 0:
			for i in range(ss - S + 1):
				res += d[i] * d[i + S]
		if S == 0:
			for i in range(ss + 1):
				res += d[i] * (d[i] - 1) // 2
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		counter = collections.Counter({0: 1})
		s = 0
		c = 0
		for a in A:
			s += a
			c += counter[s - S]
			counter[s] += 1
		return c


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		p = [0]
		for i in A:
			p.append(p[-1] + i)
		cnt = collections.Counter()
		res = 0
		for i in p:
			res += cnt[i]
			cnt[i + S] += 1
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		d = defaultdict(int)
		d[0] = 1
		psum = res = 0
		for a in A:
			psum += a
			res += d[psum - S]
			d[psum] += 1
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		numToFreq = {}
		numToFreq[0] = 1
		now_sum = 0
		ans = 0
		for num in A:
			now_sum += num
			if now_sum == S:
				ans += numToFreq[0]
			elif now_sum > S:
				if now_sum - S in numToFreq:
					ans += numToFreq[now_sum - S]
			if now_sum not in numToFreq:
				numToFreq[now_sum] = 0
			numToFreq[now_sum] += 1
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		n = len(A)

		def num_subar_at_most(lim):
			subarr_ct = 0
			r_sum = 0
			li = 0
			for ri in range(n):
				r_sum += A[ri]
				while li <= ri and r_sum > lim:
					r_sum -= A[li]
					li += 1
				subarr_ct += ri - li + 1
			return subarr_ct
		return num_subar_at_most(S) - num_subar_at_most(S - 1)


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		cumsum = 0
		dic = {0: [-1]}
		ans = 0
		for (i, num) in enumerate(A):
			cumsum += num
			dic[cumsum] = dic.get(cumsum, []) + [i]
			ans += len(dic.get(cumsum - S, []))
		return ans - (0 if S else len(A))


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		ans = 0
		(il, ih, sl, sh) = (0, 0, 0, 0)
		for (j, x) in enumerate(A):
			sl += x
			while il < j and sl > S:
				sl -= A[il]
				il += 1
			sh += x
			while ih < j and (sh > S or (sh == S and (not A[ih]))):
				sh -= A[ih]
				ih += 1
			if sl == S:
				ans += ih - il + 1
		return ans


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:

		def helper(k):
			if k < 0:
				return 0
			num = 0
			i = 0
			for j in range(len(A)):
				k -= A[j]
				while k < 0:
					k += A[i]
					i += 1
				num += j - i + 1
			return num
		return helper(S) - helper(S - 1)


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		n = len(A)
		(i, j, total, res) = (0, 0, 0, 0)
		while j < n:
			total += A[j]
			while i < j and total > S:
				total -= A[i]
				i += 1
			if total == S:
				res += 1
				k = i
				while k < j and A[k] == 0:
					k += 1
					res += 1
			j += 1
		return res


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		count = {0: 1}
		total = maxS = 0
		for num in A:
			total += num
			maxS += count.get(total - S, 0)
			count[total] = count.get(total, 0) + 1
		return maxS


class Solution:

	def numSubarraysWithSum(self, A: List[int], S: int) -> int:
		ctr = collections.Counter({0: 1})
		prefix = res = 0
		for a in A:
			prefix += a
			res += ctr[prefix - S]
			ctr[prefix] += 1
		return res