Given an array A[] of N integers and an integer M. The task is to count the number of subarrays which have median M.
Median of an array is defined as below:
1. If N is odd, the median value is the number that is in the middle after sorting the array.
2. if N is even, the median value is the left of the two middle numbers after sorting the array. 
Example 1:
Input:
N = 5, M = 2
A[] = {2, 1, 3, 5, 4}
Output: 
3
Explanation: 
The subarrays which has median equal to M
are [2], [2,1,3] and [2,1,3,5]
Example 2:
Input:
N = 1, M = 1
A[] = {1}
Output: 
1
Explanation: 
The subarrays which has median equal to M
is [1].
Your Task: 
You don't need to read input or print anything. Complete the function countSubarray( ) which takes the integer N , the array A[], and the integer M as input parameters and returns the number of subarays. 
Expected Time Complexity: O(N)
Expected Auxiliary Space: O(N)
Constraints:
1 ≤ N ≤ 10^{5}
1 ≤ A[] ≤ 10^{5}
1 ≤ M ≤ N

Source: geeksforgeeks | Difficulty: MEDIUM_HARD
Tags: Sorting

Data structures

Starter code: #User function Template for python3

class Solution:

    def countSubarray(self, N, A, M): 

        #code here



Solutions:
class Solution:

	def countSubarray(self, N, A, M):

		def fn(n, A, m):
			mp = [0] * (2 * n + 1)
			(cur, tot, ans) = (n, 0, 0)
			mp[cur] += 1
			for i in range(n):
				x = -1
				if A[i] >= m:
					x = 1
				if x == -1:
					tot -= mp[cur + x]
				else:
					tot += mp[cur]
				cur += x
				ans += tot
				mp[cur] += 1
			return ans
		return fn(N, A, M) - fn(N, A, M + 1)


class Solution:

	def countSubarray(self, n, nums, k):

		def rohit(n, nums, k):
			d = {0: 1}
			ans = 0
			cur = 0
			sum = 0
			n = len(nums)
			for i in range(n):
				if nums[i] < k:
					sum -= 1
					if sum in d:
						cur -= d[sum]
				else:
					if sum in d:
						cur += d[sum]
					sum += 1
				ans += cur
				if sum in d:
					d[sum] += 1
				else:
					d[sum] = 1
			return ans
		return rohit(n, nums, k) - rohit(n, nums, k + 1)


class Solution:

	def countSubarray(self, N, A, M):

		def greaterMedian(m):
			tot = 0
			ans = 0
			pref = [0] * (2 * N + 1)
			cur = N
			pref[cur] += 1
			for i in A:
				x = -1
				if i >= m:
					x = 1
				if x == 1:
					tot += pref[cur]
				else:
					tot -= pref[cur + x]
				ans += tot
				cur += x
				pref[cur] += 1
			return ans
		return greaterMedian(M) - greaterMedian(M + 1)


def count(N, nums, k):
	d = {0: 1}
	ans = 0
	cur = 0
	sum = 0
	for i in range(N):
		if nums[i] < k:
			sum -= 1
			if sum in d:
				cur -= d[sum]
		else:
			if sum in d:
				cur += d[sum]
			sum += 1
		ans += cur
		if sum in d:
			d[sum] += 1
		else:
			d[sum] = 1
	return ans

class Solution:

	def countSubarray(self, N, A, M):
		a = count(N, A, M)
		b = count(N, A, M + 1)
		return a - b


class Solution:

	def countSubarray(self, N, A, M):

		def rohit(N, A, M):
			d = {0: 1}
			ans = 0
			cur = 0
			sum = 0
			N = len(A)
			for i in range(N):
				if A[i] < M:
					sum -= 1
					if sum in d:
						cur -= d[sum]
				else:
					if sum in d:
						cur += d[sum]
					sum += 1
				ans += cur
				if sum in d:
					d[sum] += 1
				else:
					d[sum] = 1
			return ans
		return rohit(N, A, M) - rohit(N, A, M + 1)


class Solution:

	def helper(self, N, A, M):
		deltTracker = dict()
		deltTracker[0] = 1
		(out, local, delta) = (0, 0, 0)
		for i in A:
			j = 1 if i >= M else -1
			local = local + deltTracker.get(delta, 0) if j > 0 else local - deltTracker.get(delta - 1, 0)
			delta += j
			deltTracker[delta] = deltTracker.get(delta, 0) + 1
			out += local
		return out

	def countSubarray(self, N, A, M):
		return self.helper(N, A, M) - self.helper(N, A, M + 1)


class Solution:

	def c(self, N, A, M):
		mp = {}
		mp[0] = 1
		ans = 0
		ind = 0
		ele = 0
		for i in A:
			if i >= M:
				ele = ele + mp[ind]
				ind = ind + 1
			else:
				ind = ind - 1
				if ind in mp:
					ele = ele - mp[ind]
				else:
					mp[ind] = 0
					ele = ele - mp[ind]
			ans = ans + ele
			if ind in mp:
				mp[ind] += 1
			else:
				mp[ind] = 1
		return ans

	def countSubarray(self, N, A, M):
		return self.c(N, A, M) - self.c(N, A, M + 1)


class Solution:

	def countSubarray(self, N, A, M):

		def count(m):
			ans = 0
			s = 0
			cnts = {0: 1}
			cursum = 0
			for (i, x) in enumerate(A):
				if x >= m:
					b = 1
					cursum += cnts.get(s, 0)
				else:
					b = -1
					cursum -= cnts.get(s - 1, 0)
				s += b
				ans += cursum
				if s not in cnts:
					cnts[s] = 0
				cnts[s] += 1
			return ans
		return count(M) - count(M + 1)


from collections import Counter

class Solution:

	def countSubarray(self, N, A, M):

		def count(M):
			res = x = 0
			s = 0
			d = Counter()
			d[0] = 1
			for i in range(N):
				if A[i] >= M:
					x += d[s]
					s += 1
				else:
					s -= 1
					x -= d[s]
				d[s] += 1
				res += x
			return res
		return count(M) - count(M + 1)


class Solution:

	def countSubarray(self, N, A, M):

		def count(M):
			res = x = 0
			s = 0
			d = {}
			d[0] = 1
			for i in range(N):
				if A[i] >= M:
					x += d.get(s, 0)
					s += 1
				else:
					s -= 1
					x -= d.get(s, 0)
				d[s] = d.get(s, 0) + 1
				res += x
			return res
		return count(M) - count(M + 1)


class Solution:

	def solve(self, A, N, M):
		map = {}
		cur = 0
		tot = 0
		ans = 0
		map[cur] = 1
		for i in range(N):
			x = -1
			if A[i] >= M:
				x = 1
			if x == -1:
				if cur - 1 in map:
					tot -= map[cur - 1]
				else:
					map[cur - 1] = 0
			else:
				tot += map[cur]
			cur += x
			ans += tot
			if cur in map:
				map[cur] = map[cur] + 1
			else:
				map[cur] = 1
		return ans

	def countSubarray(self, N, A, M):
		return self.solve(A, N, M) - self.solve(A, N, M + 1)

	def countSubarray_bf(self, N, A, M):
		if N == 1:
			return 1
		sa_cnt = 0
		for j in range(N):
			for k in range(1, N - j + 1):
				sa = A[j:j + k]
				sa.sort()
				if k % 2:
					if sa[k // 2] == M:
						sa_cnt += 1
				elif sa[k // 2 - 1] == M:
					sa_cnt += 1
		return sa_cnt


class Solution:

	def countSubarray(self, N, A, M):

		def fn(M):
			cur = 0
			total = 0
			res = 0
			mem = {0: 1}
			for i in range(N):
				if A[i] >= M:
					x = 1
				else:
					x = -1
				if x == -1:
					total -= mem.get(cur - 1, 0)
				else:
					total += mem.get(cur, 0)
				cur += x
				res += total
				mem[cur] = 1 + mem.get(cur, 0)
			return res
		return fn(M) - fn(M + 1)


class Solution:

	def countSubarray(self, N, A, M):

		def abc(n, A, m):
			l = [0] * (2 * n + 1)
			(c, t, a) = (n, 0, 0)
			l[c] += 1
			for i in range(n):
				x = -1
				if A[i] >= m:
					x = 1
				if x == -1:
					t -= l[c + x]
				else:
					t += l[c]
				c += x
				a += t
				l[c] += 1
			return a
		return abc(N, A, M) - abc(N, A, M + 1)


class Solution:

	def countSubarray(self, N, A, M):

		def feasible(n, A, m):
			mp = [0] * (2 * n + 1)
			(curr, total, res) = (n, 0, 0)
			mp[curr] += 1
			for i in range(n):
				x = -1
				if A[i] >= m:
					x = 1
				if x == -1:
					total -= mp[curr + x]
				else:
					total += mp[curr]
				curr += x
				res += total
				mp[curr] += 1
			return res
		return feasible(N, A, M) - feasible(N, A, M + 1)


class Solution:

	def solve(self, N, A, M):
		mp = [0 for _ in range(2 * N + 1)]
		(curr, tot, ans) = (N, 0, 0)
		mp[curr] += 1
		for i in range(N):
			x = -1
			if A[i] >= M:
				x = 1
			if x == -1:
				tot -= mp[curr + x]
			else:
				tot += mp[curr]
			curr += x
			ans += tot
			mp[curr] += 1
		return ans

	def countSubarray(self, N, A, M):
		return self.solve(N, A, M) - self.solve(N, A, M + 1)


class Solution:

	def countSubarray(self, N, A, M):
		a = self.r(N, A, M)
		b = self.r(N, A, M + 1)
		return a - b

	def r(self, n, a, m):
		l = [0] * (2 * n + 1)
		c = n
		y = 0
		s = 0
		l[c] += 1
		for i in range(n):
			x = -1
			if a[i] >= m:
				x = 1
			if x == -1:
				y = y - l[c + x]
			else:
				y += l[c]
			c += x
			s += y
			l[c] += 1
		return s


class Solution:

	def countSubarray(self, N, A, M):

		def func(n, A, m):
			mp = [0] * (2 * n + 1)
			(current, total, answer) = (n, 0, 0)
			mp[current] += 1
			for i in range(n):
				z = -1
				if A[i] >= m:
					z = 1
				if z == -1:
					total -= mp[current + z]
				else:
					total += mp[current]
				current += z
				answer += total
				mp[current] += 1
			return answer
		return func(N, A, M) - func(N, A, M + 1)


class Solution:

	def countSubarray(self, N, A, M):
		return self.getsub(N, A, M) - self.getsub(N, A, M + 1)

	def getsub(self, n, a, m):
		li = [0] * 2 * (n + 1)
		ans = 0
		cur = n
		tot = 0
		li[cur] += 1
		for i in range(n):
			if a[i] >= m:
				tot += li[cur]
				cur += 1
			else:
				tot -= li[cur - 1]
				cur -= 1
			ans += tot
			li[cur] += 1
		return ans


class Solution:

	def solve(self, A, N, M):
		map = {}
		cur = 0
		tot = 0
		ans = 0
		map[cur] = 1
		for i in range(N):
			if A[i] >= M:
				x = 1
			else:
				x = -1
			if x == -1:
				tot -= map.get(cur - 1, 0)
			else:
				tot += map.get(cur, 0)
			cur += x
			ans += tot
			map[cur] = map.get(cur, 0) + 1
		return ans

	def countSubarray(self, N, A, M):
		v1 = self.solve(A, N, M)
		v2 = self.solve(A, N, M + 1)
		return v1 - v2


from collections import defaultdict

class Solution:

	def countSubarray(self, N, A, M):

		def getval(M):
			curr = total = sol = 0
			temp = defaultdict(int)
			temp[curr] = 1
			for ele in A:
				if ele >= M:
					val = 1
				else:
					val = -1
				if val == -1:
					total -= temp[curr - 1]
				else:
					total += temp[curr]
				curr += val
				sol += total
				temp[curr] += 1
			return sol
		return getval(M) - getval(M + 1)