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

import sys


def build(tree, A, start, end, node):
    if (start == end):
        tree[node] = A[start]
        return tree[node]
    mid = (int)((start + end) / 2)
    X = build(tree, A, start, mid, 2 * node + 1)
    Y = build(tree, A, mid + 1, end, 2 * node + 2)
    return (tree[node] == min(X, Y))


def query(tree, start, end, l, r, node):
    if (start > r or end < l):
        return sys .maxsize
    if (start >= l and end <= r):
        return tree[node]
    mid = (int)((start + end) / 2)
    X = query(tree, start, mid, l, r, 2 * node + 1)
    Y = query(tree, mid + 1, end, l, r, 2 * node + 2)
    return min(X, Y)


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]
    tree = [0] * (4 * N + 5)
    build(tree, A, 0, N - 1, 0)
    for i in range(N):
        start = i
        end = N - 1
        max_index = i
        while (start <= end):
            mid = (int)((start + end) / 2)
            min_element = query(tree, 0, N - 1, i, mid, 0)
            expected_sum = (mid - i + 1) * min_element
            actual_sum = preSum[mid + 1] - preSum[i]
            if (actual_sum - expected_sum <= K):
                start = mid + 1
                max_index = max(max_index, mid)
            else:
                end = mid - 1
        res = max(res, max_index - i + 2)
    return res


arr = [1, 7, 3, 4, 5, 6]
k = 6
n = 6
print(longestSubArray(arr, n, k))

# Expected output:
4