Count of subarrays having sum equal to its length | Python3 program for the above approach ; Function that counts the subarrays with sum of its elements as its length ; Decrementing all the elements of the array by 1 ; Making prefix sum array ; Declare map to store count of elements upto current element ; To count all the subarrays whose prefix sum is 0 ; Iterate the array ; Increment answer by count of current element of prefix array ; Return the answer ; Given array arr [ ] ; Function call
from collections import defaultdict NEW_LINE def countOfSubarray ( arr , N ) : NEW_LINE INDENT for i in range ( N ) : NEW_LINE INDENT arr [ i ] -= 1 NEW_LINE DEDENT pref = [ 0 ] * N NEW_LINE pref [ 0 ] = arr [ 0 ] NEW_LINE for i in range ( 1 , N ) : NEW_LINE INDENT pref [ i ] = pref [ i - 1 ] + arr [ i ] NEW_LINE DEDENT mp = defaultdict ( lambda : 0 ) NEW_LINE answer = 0 NEW_LINE mp [ 0 ] += 1 NEW_LINE for i in range ( N ) : NEW_LINE INDENT answer += mp [ pref [ i ] ] NEW_LINE mp [ pref [ i ] ] += 1 NEW_LINE DEDENT return answer NEW_LINE DEDENT arr = [ 1 , 1 , 0 ] NEW_LINE N = len ( arr ) NEW_LINE print ( countOfSubarray ( arr , N ) ) NEW_LINE