Count of subarrays having sum equal to its length | Set 2 | Python3 program for the above approach ; Function that counts the subarrays with sum of its elements as its length ; Store count of elements upto current element with length i ; Stores the final count of subarray ; Stores the prefix sum ; If size of subarray is 1 ; Iterate the array ; Find the sum ; Update frequency in map ; Print the total count ; Driver code ; Given array ; Size of the array ; Function Call
from collections import defaultdict NEW_LINE def countOfSubarray ( arr , N ) : NEW_LINE INDENT mp = defaultdict ( lambda : 0 ) NEW_LINE answer = 0 NEW_LINE sum = 0 NEW_LINE mp [ 1 ] += 1 NEW_LINE for i in range ( N ) : NEW_LINE INDENT sum += arr [ i ] NEW_LINE answer += mp [ sum - i ] NEW_LINE mp [ sum - i ] += 1 NEW_LINE DEDENT print ( answer ) NEW_LINE DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT arr = [ 1 , 0 , 2 , 1 , 2 , - 2 , 2 , 4 ] NEW_LINE N = len ( arr ) NEW_LINE countOfSubarray ( arr , N ) NEW_LINE DEDENT