Maximize length of longest subarray consisting of same elements by at most K decrements | Python3 program to implement the above approach ; Function to construct Segment Tree to return the minimum element in a range ; If leaf nodes of the tree are found ; Update the value in segment tree from given array ; Divide left and right subtree ; Stores smallest element in subarray : arr [ start ] , arr [ mid ] ; Stores smallest element in subarray : arr [ mid + 1 ] , arr [ end ] ; Stores smallest element in subarray : arr [ start ] , arr [ end ] ; Function to find the smallest element present in a subarray ; If elements of the subarray are not in the range [ l , r ] ; If all the elements of the subarray are in the range [ l , r ] ; Divide tree into left and right subtree ; Stores smallest element in left subtree ; Stores smallest element in right subtree ; Function that find length of longest subarray with all equal elements in atmost K decrements ; Stores length of longest subarray with all equal elements in atmost K decrements . ; Store the prefix sum array ; Calculate the prefix sum array ; Build the segment tree for range min query ; Traverse the array ; Stores start index of the subarray ; Stores end index of the subarray ; Stores end index of the longest subarray ; Performing the binary search to find the endpoint for the selected range ; Find the mid for binary search ; Find the smallest element in range [ i , mid ] using Segment Tree ; Stores total sum of subarray after K decrements ; Stores sum of elements of subarray before K decrements ; If subarray found with all equal elements ; Update start ; Update max_index ; If false , it means that the selected range is invalid ; Update end ; Store the length of longest subarray ; Return result ; Driver Code
import sys NEW_LINE def build ( tree , A , start , end , node ) : NEW_LINE INDENT if ( start == end ) : NEW_LINE INDENT tree [ node ] = A [ start ] NEW_LINE return tree [ node ] NEW_LINE DEDENT mid = ( int ) ( ( start + end ) / 2 ) NEW_LINE X = build ( tree , A , start , mid , 2 * node + 1 ) NEW_LINE Y = build ( tree , A , mid + 1 , end , 2 * node + 2 ) NEW_LINE return ( tree [ node ] == min ( X , Y ) ) NEW_LINE DEDENT def query ( tree , start , end , l , r , node ) : NEW_LINE INDENT if ( start > r or end < l ) : NEW_LINE INDENT return sys . maxsize NEW_LINE DEDENT if ( start >= l and end <= r ) : NEW_LINE INDENT return tree [ node ] NEW_LINE DEDENT mid = ( int ) ( ( start + end ) / 2 ) NEW_LINE X = query ( tree , start , mid , l , r , 2 * node + 1 ) NEW_LINE Y = query ( tree , mid + 1 , end , l , r , 2 * node + 2 ) NEW_LINE return min ( X , Y ) NEW_LINE DEDENT def longestSubArray ( A , N , K ) : NEW_LINE INDENT res = 1 NEW_LINE preSum = [ 0 ] * ( N + 1 ) NEW_LINE preSum [ 0 ] = A [ 0 ] NEW_LINE for i in range ( N ) : NEW_LINE INDENT preSum [ i + 1 ] = preSum [ i ] + A [ i ] NEW_LINE DEDENT tree = [ 0 ] * ( 4 * N + 5 ) NEW_LINE build ( tree , A , 0 , N - 1 , 0 ) NEW_LINE for i in range ( N ) : NEW_LINE INDENT start = i NEW_LINE end = N - 1 NEW_LINE max_index = i NEW_LINE while ( start <= end ) : NEW_LINE INDENT mid = ( int ) ( ( start + end ) / 2 ) NEW_LINE min_element = query ( tree , 0 , N - 1 , i , mid , 0 ) NEW_LINE expected_sum = ( mid - i + 1 ) * min_element NEW_LINE actual_sum = preSum [ mid + 1 ] - preSum [ i ] NEW_LINE if ( actual_sum - expected_sum <= K ) : NEW_LINE INDENT start = mid + 1 NEW_LINE max_index = max ( max_index , mid ) NEW_LINE DEDENT else : NEW_LINE INDENT end = mid - 1 NEW_LINE DEDENT DEDENT res = max ( res , max_index - i + 2 ) NEW_LINE DEDENT return res NEW_LINE DEDENT arr = [ 1 , 7 , 3 , 4 , 5 , 6 ] NEW_LINE k = 6 NEW_LINE n = 6 NEW_LINE print ( longestSubArray ( arr , n , k ) ) NEW_LINE