Number of subarrays having sum of the form k ^ m , m >= 0 | Python3 implementation of the above approach ; Function to count number of sub - arrays whose sum is k ^ p where p >= 0 ; If m [ a + b ] = c , then add c to the current sum . ; Increase count of prefix sum . ; If m [ a + b ] = c , then add c to the current sum . ; Increase count of prefix sum . ; b = k ^ p , p >= 0 ; k ^ m can be maximum equal to 10 ^ 14. ; If m [ a + b ] = c , then add c to the current sum . ; Increase count of prefix sum . ; Driver code
from collections import defaultdict NEW_LINE MAX = 100005 NEW_LINE def partial_sum ( prefix_sum , arr , n ) : NEW_LINE INDENT for i in range ( 1 , n + 1 ) : NEW_LINE INDENT prefix_sum [ i ] = ( prefix_sum [ i - 1 ] + arr [ i - 1 ] ) NEW_LINE DEDENT return prefix_sum NEW_LINE DEDENT def countSubarrays ( arr , n , k ) : NEW_LINE INDENT prefix_sum = [ 0 ] * MAX NEW_LINE prefix_sum [ 0 ] = 0 NEW_LINE prefix_sum = partial_sum ( prefix_sum , arr , n ) NEW_LINE if ( k == 1 ) : NEW_LINE INDENT sum = 0 NEW_LINE m = defaultdict ( int ) NEW_LINE for i in range ( n , - 1 , - 1 ) : NEW_LINE INDENT if ( ( prefix_sum [ i ] + 1 ) in m ) : NEW_LINE INDENT sum += m [ prefix_sum [ i ] + 1 ] NEW_LINE DEDENT m [ prefix_sum [ i ] ] += 1 NEW_LINE DEDENT return sum NEW_LINE DEDENT if ( k == - 1 ) : NEW_LINE INDENT sum = 0 NEW_LINE m = defaultdict ( int ) NEW_LINE for i in range ( n , - 1 , - 1 ) : NEW_LINE INDENT if ( ( prefix_sum [ i ] + 1 ) in m ) : NEW_LINE INDENT sum += m [ prefix_sum [ i ] + 1 ] NEW_LINE DEDENT if ( ( prefix_sum [ i ] - 1 ) in m ) : NEW_LINE INDENT sum += m [ prefix_sum [ i ] - 1 ] NEW_LINE DEDENT m [ prefix_sum [ i ] ] += 1 NEW_LINE DEDENT return sum NEW_LINE DEDENT sum = 0 NEW_LINE m = defaultdict ( int ) NEW_LINE for i in range ( n , - 1 , - 1 ) : NEW_LINE INDENT b = 1 NEW_LINE while ( True ) : NEW_LINE INDENT if ( b > 100000000000000 ) : NEW_LINE INDENT break NEW_LINE DEDENT if ( ( prefix_sum [ i ] + b ) in m ) : NEW_LINE INDENT sum += m [ prefix_sum [ i ] + b ] NEW_LINE DEDENT b *= k NEW_LINE DEDENT m [ prefix_sum [ i ] ] += 1 NEW_LINE DEDENT return sum NEW_LINE DEDENT if __name__ == " _ _ main _ _ " : NEW_LINE INDENT arr = [ 2 , 2 , 2 , 2 ] NEW_LINE n = len ( arr ) NEW_LINE k = 2 NEW_LINE print ( countSubarrays ( arr , n , k ) ) NEW_LINE DEDENT