2333. Minimum Sum of Squared Difference

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.

class Solution {
public:
    long long minSumSquareDiff(vector<int>& A, vector<int>& B, int k1, int k2) {
        long long res = 0, n = A.size();
        long long diffsum = 0, k = k1*1ll + k2;
        unordered_map<long long, long long> freq;
        for(int i = 0; i < n; i++) {
            freq[abs(A[i] - B[i])]++;
        }
        vector<pair<long long, long long>> C;
        for(auto [k,v] : freq)
            C.push_back({k,v});
        sort(begin(C), end(C));
        while(k and !C.empty()) {
            auto [diff, count] = C.back(); C.pop_back();
            auto diff2 = C.size() == 0 ? 0ll : C.back().first;
            if((diff - diff2) * count <= k) {
                k -= (diff - diff2) * count;
                if(C.empty()) return 0;
                C.back().second += count;
            } else {
                int minus = k / count, remain = k % count;
                k = 0;
                C.push_back({diff-minus, count-remain});
                C.push_back({diff-minus-1,remain});
            }
        }
        for(auto [diff, count] : C)
            res += diff * diff * count;
        return res;
    }
};