2333. Minimum Sum of Squared Difference

/**
 * Problem: LeetCode 2333 - Minimum Sum of Squared Difference
 * Difficulty: Medium
 * Language: Java
 * Author: CodeByTushu
 */
class Solution {
    public long minSumSquareDiff(int[] nums1, int[] nums2, int k1, int k2) {
        int n = nums1.length;
        int maxDiff = 100000;
        long[] cnt = new long[maxDiff + 2];
        for (int i = 0; i < n; i++) {
            cnt[Math.abs(nums1[i] - nums2[i])]++;
        }
        long k = (long) k1 + k2;
        for (int d = maxDiff; d >= 1 && k > 0; d--) {
            if (cnt[d] == 0) continue;
            if (cnt[d] <= k) {
                k -= cnt[d];
                cnt[d - 1] += cnt[d];
                cnt[d] = 0;
            } else {
                cnt[d - 1] += k;
                cnt[d] -= k;
                k = 0;
            }
        }
        long ans = 0;
        for (int d = 1; d <= maxDiff; d++) {
            ans += (long) d * d * cnt[d];
        }
        return ans;
    }
}

"""
Problem: LeetCode 2333 - Minimum Sum of Squared Difference
Difficulty: Medium
Language: Python
Author: CodeByTushu
"""
class Solution:
    def minSumSquareDiff(self, nums1, nums2, k1, k2):
        max_diff = 100000
        cnt = [0] * (max_diff + 2)

        for a, b in zip(nums1, nums2):
            cnt[abs(a - b)] += 1

        k = k1 + k2

        for d in range(max_diff, 0, -1):
            if k == 0:
                break

            if cnt[d] <= k:
                k -= cnt[d]
                cnt[d - 1] += cnt[d]
                cnt[d] = 0
            else:
                cnt[d - 1] += k
                cnt[d] -= k
                k = 0

        ans = 0
        for d in range(1, max_diff + 1):
            ans += d * d * cnt[d]

        return ans

/**
 * Problem: LeetCode 2333 - Minimum Sum of Squared Difference
 * Difficulty: Medium
 * Language: C++
 * Author: CodeByTushu
 */
#include <vector>
#include <cmath>
#include <cstdlib>

using namespace std;

/**
 * Problem: LeetCode 2333 - Minimum Sum of Squared Difference
 * Difficulty: Medium
 * Language: C++
 * Author: CodeByTushu
 */
class Solution {
public:
    long long minSumSquareDiff(vector<int>& nums1, vector<int>& nums2,
                               int k1, int k2) {
        const int maxDiff = 100000;
        vector<long long> cnt(maxDiff + 2, 0);

        int n = nums1.size();

        for (int i = 0; i < n; i++) {
            cnt[abs(nums1[i] - nums2[i])]++;
        }

        long long k = (long long)k1 + k2;

        for (int d = maxDiff; d >= 1 && k > 0; d--) {
            if (cnt[d] == 0) continue;

            if (cnt[d] <= k) {
                k -= cnt[d];
                cnt[d - 1] += cnt[d];
                cnt[d] = 0;
            } else {
                cnt[d - 1] += k;
                cnt[d] -= k;
                k = 0;
            }
        }

        long long ans = 0;

        for (int d = 1; d <= maxDiff; d++) {
            ans += 1LL * d * d * cnt[d];
        }

        return ans;
    }
};

/**
 * Problem: LeetCode 2333 - Minimum Sum of Squared Difference
 * Difficulty: Medium
 * Language: JavaScript
 * Author: CodeByTushu
 */
class Solution {
    minSumSquareDiff(nums1, nums2, k1, k2) {
        const maxDiff = 100000;
        const cnt = new Array(maxDiff + 2).fill(0);

        for (let i = 0; i < nums1.length; i++) {
            cnt[Math.abs(nums1[i] - nums2[i])]++;
        }

        let k = k1 + k2;

        for (let d = maxDiff; d >= 1 && k > 0; d--) {
            if (cnt[d] === 0) continue;

            if (cnt[d] <= k) {
                k -= cnt[d];
                cnt[d - 1] += cnt[d];
                cnt[d] = 0;
            } else {
                cnt[d - 1] += k;
                cnt[d] -= k;
                k = 0;
            }
        }

        let ans = 0;

        for (let d = 1; d <= maxDiff; d++) {
            ans += d * d * cnt[d];
        }

        return ans;
    }
}