Given a binary string S and an integer K, the task is to count the number of substrings that contain exactly K ones.
Example 1:
Input : S = 10010, K = 1
Output : 9
Explanation: The 9 substrings containing 
one 1 are, 1, 10, 100, 001, 01,
1, 10, 0010 and 010.
Example 2:
Input : S = 111, K = 2 
Output : 2 
Explanation: The 2 substrings containing
two 1 are 11, 11.
Your Task:  
You don't need to read input or print anything. Your task is to complete the function countOfSubstringWithKOnes() which takes the string S and an integer K as inputs and returns the count.
Expected Time Complexity: O(|S|)
Expected Auxiliary Space: O(|S|)
Constraints:
1 ≤ N ≤ 10^{5}
1 ≤ K ≤ 10^{4}

Source: geeksforgeeks | Difficulty: MEDIUM
Tags: String algorithms

Data structures

Starter code: #User function Template for python3



class Solution:

    def countOfSubstringWithKOnes(self, S, K):

        # code here

Solutions:
class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		cm = 0
		mp = dict()
		n = len(S)
		ans = 0
		for i in range(n):
			cm += ord(S[i]) - 48
			if cm == k:
				ans += 1
			if cm - k in mp:
				ans += mp[cm - k]
			if cm not in mp:
				mp[cm] = 1
			else:
				mp[cm] += 1
		return ans


class Solution:

	def solve(self, s, k):
		ans = 0
		n = len(s)
		cnt = 0
		j = 0
		for i in range(n):
			if s[i] == '1':
				cnt += 1
			while j < n and cnt > k:
				if s[j] == '1':
					cnt -= 1
				j += 1
			ans += i - j + 1
		return ans

	def countOfSubstringWithKOnes(self, S, K):
		if K == 0:
			return self.solve(S, K)
		return self.solve(S, K) - self.solve(S, K - 1)


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		freq = {}
		S = list(str(S))
		presum = 0
		freq[0] = 1
		c = 0
		for i in S:
			presum += int(i)
			if presum - K in freq:
				c += freq[presum - K]
			if presum not in freq:
				freq[presum] = 0
			freq[presum] += 1
		return c


class Solution:

	def countOfSubstringWithKOnes(self, S, k):
		if k == 0:
			count = 0
			c = 0
			for i in S:
				if i == '0':
					c += 1
				else:
					count += c * (c + 1) // 2
					c = 0
			if c:
				count += c * (c + 1) // 2
			return count
		nums = list(S)
		count = 0
		(a, b, c) = (0, 0, 1)
		ln = len(nums)
		for i in range(ln):
			nums[i] = nums[i] == '1'
		while a < ln and nums[a] == 0:
			c += 1
			a += 1
		for i in range(ln):
			k -= nums[i]
			if k == 0:
				count += c
			elif k == -1:
				k = 0
				a += 1
				c = 1
				while nums[a] == 0:
					c += 1
					a += 1
				count += c
		return count


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		count = 0
		n = len(s)
		feq = [0 for i in range(n + 1)]
		feq[0] = 1
		res = 0
		for i in range(n):
			count += ord(s[i]) - ord('0')
			if count >= K:
				res += feq[count - k]
			feq[count] += 1
		return res


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		count = 0
		S = str(' '.join(S))
		import collections
		mapp = collections.defaultdict(int)
		mapp[0] = 1
		summ = 0
		l = list(map(int, S.split(' ')))
		for ele in l:
			summ += ele
			count += mapp[summ - K]
			mapp[summ] += 1
		return count


class Solution:

	def countOfSubstringWithKOnes(self, s, k):
		cum_sum = 0
		mp = {}
		cnt = 0
		arr = [int(i) for i in s]
		mp[0] = [-1]
		n = len(s)
		for i in range(n):
			cum_sum = cum_sum + arr[i]
			if cum_sum - k in mp:
				cnt = cnt + len(mp[cum_sum - k])
			mp.setdefault(cum_sum, []).append(i)
		return cnt


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		n = len(S)
		dp = [0] * (n + 1)
		dp[0] = 1
		one_count = 0
		res = 0
		for c in S:
			if c == '1':
				one_count += 1
			if one_count >= K:
				res += dp[one_count - K]
			dp[one_count] += 1
		return res


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		res = 0
		freq = [0 for i in range(len(S) + 1)]
		freq[0] = 1
		countOne = 0
		for i in range(len(S)):
			countOne += ord(S[i]) - ord('0')
			if countOne >= k:
				res += freq[countOne - k]
			freq[countOne] += 1
		return res


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		d = {}
		count = 0
		Sum = 0
		for num in S:
			Sum = Sum + int(num)
			if Sum == K:
				count = count + 1
			if Sum - K in d:
				count = count + d[Sum - K]
			if Sum not in d:
				d[Sum] = 1
			else:
				d[Sum] += 1
		return count


class Solution:

	def countOfSubstringWithKOnes(self, s, k):
		(i, j) = (0, 0)
		n = len(s)
		presum = 0
		res1 = 0
		while j < n:
			if s[j] == '1':
				presum += 1
			while presum > k:
				if s[i] == '1':
					presum -= 1
				i += 1
			res1 += j - i + 1
			j += 1
		if k == 0:
			return res1
		k -= 1
		(i, j) = (0, 0)
		presum = 0
		res2 = 0
		while j < n:
			if s[j] == '1':
				presum += 1
			while presum > k:
				if s[i] == '1':
					presum -= 1
				i += 1
			res2 += j - i + 1
			j += 1
		return res1 - res2


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		N = len(s)
		res = 0
		countOfOne = 0
		freq = [0 for i in range(N + 1)]
		freq[0] = 1
		for i in range(0, N, 1):
			countOfOne += ord(s[i]) - ord('0')
			if countOfOne >= K:
				res += freq[countOfOne - K]
			freq[countOfOne] += 1
		return res


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		a = {0: 1}
		b = []
		for i in range(len(S)):
			b.append(int(S[i]))
		sum1 = 0
		count = 0
		for i in range(len(b)):
			sum1 = sum1 + b[i]
			if sum1 - k in a:
				count = count + a[sum1 - k]
			if sum1 not in a:
				a[sum1] = 1
			else:
				a[sum1] = a[sum1] + 1
		return count
if __name__ == '__main__':
	t = int(input())
	for _ in range(t):
		inp = input().strip().split(' ')
		s = inp[0]
		k = int(inp[1])
		ob = Solution()
		print(ob.countOfSubstringWithKOnes(s, k))


from collections import Counter

class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		n = len(s)
		count = Counter()
		total = 0
		ans = 0
		count[0] = 1
		for c in s:
			total += int(c)
			if total >= K:
				ans += count[total - K]
			count[total] += 1
		return ans
if __name__ == '__main__':
	t = int(input())
	for _ in range(t):
		inp = input().strip().split(' ')
		s = inp[0]
		k = int(inp[1])
		ob = Solution()
		print(ob.countOfSubstringWithKOnes(s, k))


from collections import deque

class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		ans = 0
		left = 0
		q = deque()
		n = len(S)
		if K == 0:
			S = S.split('1')
			ans = 0
			for v in S:
				ans += (len(v) + 1) * len(v) // 2
			return ans
		for (right, val) in enumerate(S):
			if val == '1':
				if len(q) == K:
					ans += (q[0] - left + 1) * (right - q[-1])
					left = q.popleft() + 1
				q.append(right)
		if len(q) == K:
			ans += (q[0] - left + 1) * (len(S) - q[-1])
		return ans
if __name__ == '__main__':
	t = int(input())
	for _ in range(t):
		inp = input().strip().split(' ')
		s = inp[0]
		k = int(inp[1])
		ob = Solution()
		print(ob.countOfSubstringWithKOnes(s, k))


class Solution:

	def countOfSubstringWithKOnes(self, s, K):
		n = len(s)
		freq = [0 for i in range(n + 1)]
		freq[0] = 1
		c1 = 0
		res = 0
		for i in range(n):
			c1 = c1 + int(s[i])
			if c1 >= k:
				res = res + freq[c1 - k]
			freq[c1] += 1
		return res
if __name__ == '__main__':
	t = int(input())
	for _ in range(t):
		inp = input().strip().split(' ')
		s = inp[0]
		k = int(inp[1])
		ob = Solution()
		print(ob.countOfSubstringWithKOnes(s, k))


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		Sum = 0
		Dict = {0: 1}
		N = len(S)
		res = 0
		for i in range(N):
			Sum += int(S[i])
			if Sum - K >= 0 and Sum - K in Dict:
				res += Dict[Sum - K]
			if Sum not in Dict:
				Dict[Sum] = 1
			else:
				Dict[Sum] += 1
		return res


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		result = 0
		cnt = 0
		freq = [0 for i in range(len(S) + 1)]
		freq[0] = 1
		for i in range(len(S)):
			cnt += ord(S[i]) - ord('0')
			if cnt >= K:
				result += freq[cnt - K]
			freq[cnt] += 1
		return result


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		d = {}
		prefix_sum = 0
		count = 0
		for i in range(len(S)):
			prefix_sum += int(S[i])
			if prefix_sum == K:
				count += 1
			if prefix_sum - K in d:
				count += d[prefix_sum - K]
			if prefix_sum not in d:
				d[prefix_sum] = 1
			else:
				d[prefix_sum] += 1
		return count


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		l = [int(x) for x in S]
		d = {0: [-1]}
		indices = []
		sum = 0
		count = 0
		for i in range(len(l)):
			sum += l[i]
			if sum - K in d:
				for x in d[sum - K]:
					count += 1
			if sum not in d:
				d[sum] = [i]
			else:
				d[sum].append(i)
		return count


class Solution:

	def countOfSubstringWithKOnes(self, S, K):

		def s1(S, K):
			l = 0
			r = 0
			count = 0
			ret = 0
			while r < len(S):
				if S[r] == '1':
					count += 1
				while count == K + 1:
					ret += 1
					if S[l] == '1':
						count -= 1
					l += 1
				r += 1
			return ret

		def s2(S, K):
			from collections import defaultdict
			freqOfCount = defaultdict(int)
			countSoFar = 0
			ret = 0
			freqOfCount[0] = 1
			for (i, c) in enumerate(S):
				if c == '1':
					countSoFar += 1
				if countSoFar >= K:
					ret += freqOfCount[countSoFar - K]
				freqOfCount[countSoFar] += 1
			return ret
		return s2(S, K)


class Solution:

	def countOfSubstringWithKOnes(self, s, k):
		map1 = {0: 1}
		ans = 0
		curr_sum = 0
		for i in range(len(s)):
			if s[i] == '1':
				curr_sum += 1
			ans += map1[curr_sum - k] if curr_sum - k in map1 else 0
			map1[curr_sum] = map1.get(curr_sum, 0) + 1
		return ans


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		freq = dict()
		sm = 0
		c = 0
		for i in S:
			if i == '1':
				sm += 1
			if sm == K:
				c = c + 1
			if sm - K in freq:
				c = c + freq[sm - K]
			freq[sm] = freq.get(sm, 0) + 1
		return c


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		string = S
		n = K
		prefixCount = []
		freq = {}
		curCount = 0
		if n == 0:
			mainSu = 0
			curCount = 0
			for x in string:
				if x == '1':
					mainSu += curCount * (curCount + 1) // 2
					curCount = 0
				else:
					curCount += 1
			mainSu += curCount * (curCount + 1) // 2
			return mainSu
		for x in string:
			if x == '1':
				curCount += 1
			prefixCount.append(curCount)
		for x in prefixCount:
			temp = freq.get(x)
			if temp:
				freq[x] += 1
			else:
				freq[x] = 1
		mainSu = 0
		prevCount = 0
		curCount = 0
		target = n
		i = 0
		flag = True
		for x in range(len(string)):
			if string[x] == '0':
				calc = n + prefixCount[x]
				temp = freq.get(calc)
				if temp:
					mainSu += freq[calc]
			else:
				calc = n + prefixCount[x] - 1
				temp = freq.get(calc)
				if temp:
					mainSu += freq[calc]
					flag = False
		if flag:
			if n == 0:
				n = len(string)
				return n * (n + 1) // 2
			else:
				return 0
		else:
			return mainSu


class Solution:

	def countOfSubstringWithKOnes(self, S, K):
		count = 0
		ones = []
		for (index, item) in enumerate(S):
			if item == '1':
				ones.append(index)
		if len(ones) < K:
			return 0
		if K == 0:
			zeros = []
			flag = False
			count0 = 0
			for i in range(len(S)):
				if S[i] == '0':
					flag = True
					count0 += 1
				else:
					if flag:
						zeros.append(count0)
					flag = False
					count0 = 0
			zeros.append(count0)
			for num in zeros:
				count += num * (num + 1) // 2
			return int(count)
		left = 0
		next1 = K - 1
		start = ones[K - 1] if K - 1 >= 0 else 0
		for i in range(start, len(S)):
			x = -1 if left == 0 else ones[left - 1]
			if next1 < len(ones) and ones[next1] == i:
				incr = ones[left] - x
				count += incr
				left += 1
				next1 += 1
			else:
				count += incr
		return count