Maximize the length of a subarray of equal elements by performing at most K increment operations

A subarray is a contiguous part of an array i.e. it can be considered as an array inside another array. For example, take the following array,

array[] = {1, 2, 3, 4, 5, 6}

For the above array, one possible subarray is

subarry[] = {2, 3, 4}

Problem Statement

Given an array arr[] having N positive integers and a positive integer K representing the maximum number that can be added to the elements of the array. The task is to increment the elements of the array by most K increment operations and return the maximum possible subarray having equal elements.

Sample Example 1

Input: arr[] = {1, 2, 2, 3, 3, 3, 4}
K = 2
Output: 5

Explanation

• Incrementing arr[1] by 1. Then arr[1] = 3

• Incrementing arr[2] by 1. Then arr[2] = 3

• Thus, the final array is {1, 3, 3, 3, 3, 3, 4}.

In the above array, the longest subarray with equal elements is of length 5.

Sample Example 2

Input: arr[] = {1, 2, 4, 5, 2}
K = 5
Output: 3

Explanation

• Incrementing arr[1] by 3. Then arr[1] = 5

• Incrementing arr[2] by 1. Then arr[2] = 5

• Thus, the final array is {1, 5, 5, 5, 2}.

In the above array, the longest subarray with equal elements is of length 3.

Approach 1: Sliding window I

For solving the above problem, the first step is to sort the array and then perform binary search to pick a value for maximum indices having the same element. Then for each picked element, use Sliding window to check if it is possible to get a subarray with all elements equal.

Exampple: C++ Implementation

In the following program, maximum length of subarray with atmost k increments are found.

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

// Function to check if subarray of length l exists having all elements equal
bool check(vector<int> vec, int l, int k, vector<int> arr){
   int i = 0;
   int j = l;
   while (j <= arr.size()){
      int maxSize = arr[j - 1];
      int total = maxSize * l;
      int curr = vec[j] - vec[i];
      if (curr + k >= total){
         return true;
      } else {
         i++;
         j++;
      }
   }
   return false;
}
// Function to find the maximum number of indices having equal elements after adding at most k numbers
int maxSubarray(vector<int> arr, int k){
   sort(arr.begin(), arr.end());
   vector<int> vec(arr.size());
   vec[1] = arr[0];
   for (int i = 1; i < vec.size() - 1; ++i){
   vec[i + 1] = vec[i] + arr[i];
   }
   int max = arr.size();
   int min = 1;
   int res = 1;
   while (min <= max){
      int mid = (max + min) / 2;
      if (check(vec, mid, k, arr)){
         res = mid;
         min = mid + 1;
      } else {
         max = mid - 1;
      }
   }
   return res;
}
int main(){
   vector<int> arr = {1, 2, 4, 5, 2};
   int k = 5;
   cout << (maxSubarray(arr, k));
   return 0;
}

Output

3

Approach 2: Sliding Window II

As a second version of using Sliding window for the above proble, we’ll take windows and for each window find the maximum element and total operations to make all elements of that window equal. If total operations are less than K then increment otherwise move rightwards. The above steps are repeated to get the longest subarray.

Example: C++ Impelementation

In the following program, maximum length of subarray with atmost k increments are found.

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

// Function to find the maximum number of indices having equal elements
int maxSubarray(int arr[], int k, int n){
   int res = 0;
   int start = 0;
   long int s = 0;
   deque<int> dq;
   for (int i = 0; i < n; i++) {
      int a = arr[i];
      while (!dq.empty() && arr[dq.front()] <= a) {
         dq.pop_front();
      }
      dq.push_back(i);
      s += a;
      long int ops = (long int)arr[dq.front()] * (res + 1) - s;
      if (ops <= (long int)k)
      res++;
      else {
         if (dq.front() == start)
         dq.pop_front();
         s -= arr[start++];
      }
   }
   return res;
}
int main(){
   int arr[] = {1, 2, 4, 5, 2};
   int k = 5;
   int n = 5;
   cout << (maxSubarray(arr, k, n));
   return 0;
}

Output

3

Conclusion

In conclusion, two approaches were discussed to solve the given problem,

The first approach uses sliding window to find subarray with equal elements. The time complexity is O(N*log(N)).

The second appraoch uses sliding window and deque data structure to find maximum length. The time complexity is O(N).