Minimize squared sum in K ops, Binary search

/*
You are given two positive 0-indexed integer arrays nums1 and nums2, both of length n.

The sum of squared difference of arrays nums1 and nums2 is defined as the sum of (nums1[i] - nums2[i])2 for each 0 <= i < n.

You are also given two positive integers k1 and k2. You can modify any of the elements of nums1 by +1 or -1 at most k1 times. Similarly, you can modify any of the elements of nums2 by +1 or -1 at most k2 times.

Return the minimum sum of squared difference after modifying array nums1 at most k1 times and modifying array nums2 at most k2 times.

Note: You are allowed to modify the array elements to become negative integers.



Example 1:

Input: nums1 = [1,2,3,4], nums2 = [2,10,20,19], k1 = 0, k2 = 0
Output: 579
Explanation: The elements in nums1 and nums2 cannot be modified because k1 = 0 and k2 = 0.
The sum of square difference will be: (1 - 2)2 + (2 - 10)2 + (3 - 20)2 + (4 - 19)2 = 579.
Example 2:

Input: nums1 = [1,4,10,12], nums2 = [5,8,6,9], k1 = 1, k2 = 1
Output: 43
Explanation: One way to obtain the minimum sum of square difference is:
- Increase nums1[0] once.
- Increase nums2[2] once.
The minimum of the sum of square difference will be:
(2 - 5)2 + (4 - 8)2 + (10 - 7)2 + (12 - 9)2 = 43.
Note that, there are other ways to obtain the minimum of the sum of square difference, but there is no way to obtain a sum smaller than 43.


Constraints:

n == nums1.length == nums2.length
1 <= n <= 105
0 <= nums1[i], nums2[i] <= 10^5
0 <= k1, k2 <= 10^9

https://leetcode.com/problems/minimum-sum-of-squared-difference/
*/

class Solution {
public:
    bool isOK(vector<long long> &diff, int K, long long guessedLevel){
        long long requiredOps=0;
        for(auto level: diff){
            if(level<=guessedLevel){ // already less or equal to guessedLevel
                continue;
            }
            requiredOps+=(level-guessedLevel);
        }
        return requiredOps<=K;  // you can bring every diff to level guessedLevel in less or equal to k1+k2 operations
    }

    long long minSumSquareDiff(vector<int>& nums1, vector<int>& nums2, int k1, int k2) {
        int n=nums1.size();
        vector<long long> diff(n);
        for(int i=0;i<n;i++){
            diff[i]=abs(nums1[i]-nums2[i]);
        }

        // Finding the level using binary search which is possible in k1+k2 ops and lead to minimum level
        long long low=0;
        long long high=INT_MAX;
        int totalOps=k1+k2;
        int minLevel=-1;
        while(low<=high){
            int mid=(low+high)/2;
            if(isOK(diff,totalOps,mid)==true){  // operations required to reach mid level is less than totalOps
                minLevel=mid;
                high=mid-1;
            }
            else{
                low=mid+1;
            }
        }

        // bring every diff to that level, and use operations for it
        for(int i=0;i<n;i++){
            if(diff[i]<=minLevel){  // already less than minLevel
                continue;
            }
            totalOps-=(diff[i]-minLevel);
            diff[i]=minLevel;
        }


        /*
           Note that it always better to reduce two largest differences by 1,
           then one of them by 2: (n - 2) * (n - 2) + n * n > (n - 1) * (n - 1) + (n - 1) * (n - 1).
        */
        if(totalOps>0){
            sort(diff.begin(),diff.end(),greater<int>());
            for(int i=0;i<n && totalOps>0;i++){
                if(diff[i]>0){
                    diff[i]--;
                    totalOps--;
                }
            }
        }

        long long res=0;
        for(auto level: diff){
            res+=pow(level,2);
        }
        return res;

    }
};

/*
THOUGHT PROCESS-

K1 K2 are not separate, you can consider them the same as K=K1+K2. Now, you want to know, if there exist any level which every difference vector diff[i] be less or equal to. This should be done using binary search, using binary search you find a level called minLevel, such that every diff[i] can be bought equal to or less than this minLevel in the given amount of operations.

Now, after you get such minLevel, its time you bring every element to that level using operations, if your diff[i] is already at that level or below it, leave it. You might have some operations left. Use them to bring big diff[i] elements one by one. This is done by sorting as you can see in the code

Then you just find the res and return
*/