HackerEarth Equal subarrays problem solution

In this HackerEarth Equal subarrays problem solution You are given an array of N non-negative integers [A1,A2,A3,…,AN], and a non-negative integer K. A subarray is an array composed of a contiguous block of original array elements. You can perform the following operation on a subarray:

Increase each element of this subarray by a non-negative number such that the total sum of all the increments does not exceed K. You must make all the elements of this subarray equal.

Determine the maximum length of a subarray in which all the elements can be made equal by performing the mentioned operation.

#include "bits/stdc++.h"
using namespace std;

int segment_tree[1000000];
long long int A[100005];    

int build(int start, int end, int node)
{
    if(start == end)
    {
        segment_tree[node] = A[start];
    }
    else
    {
        int mid = (start + end) / 2;
        segment_tree[node] = max(build(start, mid, 2*node + 1), build(mid+1, end, 2*node+2));
    }
    return segment_tree[node];
}

int query(int start, int end, int l, int r, int node)
{
    if(start > r || end < l) return -1;
    if(start >= l && end <= r) return segment_tree[node];
    int mid = (start + end) / 2;
    return max(query(start, mid, l, r, 2*node+1), query(mid+1, end, l, r, 2*node+2));
}

int main(){
    int N; cin>>N;
    int K; cin>>K;  
    for(int i=0; i<N; i++)
    {
        cin>>A[i];
    }
    long long int pref[N+1];
    pref[0] = 0;
    for(int i=0; i<N; i++)
    {
        pref[i+1] = pref[i] + A[i];
    }
    build(0, N-1, 0);
    int ans = 1;
    for(int i=0; i<N; i++)
    {
        int start = i, end = N-1, mid, max_index = i;
        while(start <= end)
        {
            mid = (start + end) / 2;
            long long int max_element = query(0, N-1, i, mid, 0);
            long long int expected_sum = (mid - i + 1) * max_element;
            long long int actual_sum = pref[mid + 1] - pref[i];
            if(expected_sum - actual_sum <= K)
            {
                start = mid + 1;
                max_index = max(max_index, mid);
            }
            else
            {
                end = mid - 1;
            }
        }
        ans = max(ans, max_index - i + 1);
    }
    cout<<ans;
    return 0;
}