Length of Longest Subarray with same elements in atmost K increments | Initialize array for segment tree ; Function that builds the segment tree to return the max element in a range ; update the value in segment tree from given array ; find the middle index ; If there are more than one elements , then recur for left and right subtrees and store the max of values in this node ; Function to return the max element in the given range ; If the range is out of bounds , return - 1 ; Function that returns length of longest subarray with same elements in atmost K increments . ; Initialize the result variable Even though the K is 0 , the required longest subarray , in that case , will also be of length 1 ; Initialize the prefix sum array ; Build the prefix sum array ; Build the segment tree for range max query ; Loop through the array with a starting point as i for the required subarray till the longest subarray is found ; Performing the binary search to find the endpoint for the selected range ; Find the mid for binary search ; Find the max element in the range [ i , mid ] using Segment Tree ; Total sum of subarray after increments ; Actual sum of elements before increments ; Check if such increment is possible If true , then current i is the actual starting point of the required longest subarray ; Now for finding the endpoint for this range Perform the Binary search again with the updated start ; Store max end point for the range to give longest subarray ; If false , it means that the selected range is invalid ; Perform the Binary Search again with the updated end ; Store the length of longest subarray ; Return result ; Driver code
segment_tree = [ 0 ] * ( 4 * 1000000 ) ; NEW_LINE def build ( A , start , end , node ) : NEW_LINE INDENT if ( start == end ) : NEW_LINE INDENT segment_tree [ node ] = A [ start ] ; NEW_LINE DEDENT else : NEW_LINE INDENT mid = ( start + end ) // 2 ; NEW_LINE segment_tree [ node ] = max ( build ( A , start , mid , 2 * node + 1 ) , build ( A , mid + 1 , end , 2 * node + 2 ) ) ; NEW_LINE DEDENT return segment_tree [ node ] ; NEW_LINE DEDENT def query ( start , end , l , r , node ) : NEW_LINE INDENT if ( start > r or end < l ) : NEW_LINE INDENT return - 1 ; NEW_LINE DEDENT if ( start >= l and end <= r ) : NEW_LINE INDENT return segment_tree [ node ] ; NEW_LINE DEDENT mid = ( start + end ) // 2 ; NEW_LINE return max ( query ( start , mid , l , r , 2 * node + 1 ) , query ( mid + 1 , end , l , r , 2 * node + 2 ) ) ; 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 build ( 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 = ( start + end ) // 2 ; NEW_LINE max_element = query ( 0 , N - 1 , i , mid , 0 ) ; NEW_LINE expected_sum = ( mid - i + 1 ) * max_element ; NEW_LINE actual_sum = preSum [ mid + 1 ] - preSum [ i ] ; NEW_LINE if ( expected_sum - actual_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 + 1 ) ; NEW_LINE DEDENT return res ; NEW_LINE DEDENT if __name__ == " _ _ main _ _ " : NEW_LINE INDENT arr = [ 2 , 0 , 4 , 6 , 7 ] ; K = 6 ; NEW_LINE N = len ( arr ) ; NEW_LINE print ( longestSubArray ( arr , N , K ) ) ; NEW_LINE DEDENT