# 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 ; 

segment_tree = [0] * (4 * 1000000)


def build(A, start, end, node):
    if (start == end):
        segment_tree[node] = A[start]
    else:
        mid = (start + end) // 2
        segment_tree[node] = max(
            build(
                A,
                start,
                mid,
                2 * node + 1),
            build(
                A,
                mid + 1,
                end,
                2 * node + 2))
    return segment_tree[node]


def query(start, end, l, r, node):
    if (start > r or end < l):
        return -1
    if (start >= l and end <= r):
        return segment_tree[node]
    mid = (start + end) // 2
    return max(
        query(
            start,
            mid,
            l,
            r,
            2 * node + 1),
        query(
            mid + 1,
            end,
            l,
            r,
            2 * node + 2))


def longestSubArray(A, N, K):
    res = 1
    preSum = [0] * (N + 1)
    preSum[0] = A[0]
    for i in range(N):
        preSum[i + 1] = preSum[i] + A[i]
    build(A, 0, N - 1, 0)
    for i in range(N):
        start = i
        end = N - 1
        max_index = i
        while (start <= end):
            mid = (start + end) // 2
            max_element = query(0, N - 1, i, mid, 0)
            expected_sum = (mid - i + 1) * max_element
            actual_sum = preSum[mid + 1] - preSum[i]
            if (expected_sum - actual_sum <= K):
                start = mid + 1
                max_index = max(max_index, mid)
            else:
                end = mid - 1
        res = max(res, max_index - i + 1)
    return res


if __name__ == "__main__":
    arr = [2, 0, 4, 6, 7]
    K = 6
    N = len(arr)
    print(longestSubArray(arr, N, K))

# Expected output:
3