human
Write me a code to find the number of sub arrays in the permutation of first N natural numbers such that their median is M and the other steps required given the following:
Find the number of sub arrays in the permutation of first N natural numbers such that their median is M | Function to return the count of sub - arrays in the given permutation of first n natural numbers such that their median is m.
If element is less than m.
If element greater than m.
If m is found.
Count the answer.
Increment Sum.
Driver code"
How can the above be solved in Python?
assistant
#Here is code to find the number of sub arrays in the permutation of first N natural numbers such that their median is M and the other steps required.
 def segments ( n , p , m ) :
	c = dict ( )
	c [ 0 ] = 1
	has = False
	Sum = 0
	ans = 0
	for r in range ( n ) :
		if ( p [ r ] < m ) :
			Sum -= 1
		elif ( p [ r ] > m ) :
			Sum += 1
		if ( p [ r ] == m ) :
			has = True
		if ( has ) :
			if ( Sum in c . keys ( ) ) :
				ans += c [ Sum ]
			if Sum - 1 in c . keys ( ) :
				ans += c [ Sum - 1 ]
		else :
			c [ Sum ] = c . get ( Sum , 0 ) + 1
	return ans
a = [ 2 , 4 , 5 , 3 , 1 ]
n = len ( a )
m = 4
print ( segments ( n , a , m ) )