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.

HackerEarth Equal subarrays problem solution.

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

Second solution

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

typedef long long ll;

const int maxn = 1e5 + 10, maxlg = 18;
int n, k, a[maxn], mx[maxlg][maxn];
ll par[maxn];

int get(int l, int r) {
    if (r <= l) return 0;
    int lg = 31 - __builtin_clz(r - l);
    return max(mx[lg][l], mx[lg][r - (1 << lg)]);
}

int main(){
    cin >> n >> k;
    for (int i = 0; i < n; ++i) {
        cin >> a[i];
        mx[0][i] = a[i];
        par[i + 1] = par[i] + a[i];
    }
    for (int i = 1; i < maxlg; ++i)
        for (int j = 0; j < n; ++j) {
            mx[i][j] = mx[i - 1][j];
            if (j + (1 << (i - 1)) < n)
                mx[i][j] = max(mx[i][j], mx[i - 1][j + (1 << (i - 1))]);
        }
    int ans = 0;
    for (int i = 0; i < n; ++i) {
        int b = i, e = n, mid;
        while (e - b > 1) {
            mid = (b + e) >> 1;
            ll val = get(i, mid + 1) * ll(mid - i + 1) - (par[mid + 1] - par[i]);
            if (val <= k)
                b = mid;
            else
                e = mid;
        }
        ans = max(ans, e - i);
    }
    cout << ans << '\n';
    return 0;
}