October 10Medium

Minimum sum of squared difference

Array
Binary Search
Greedy
Sorting
Heap

Problem

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.

Solution

class Solution:
    def minSumSquareDiff(self, nums1: list[int], nums2: list[int], k1: int, k2: int) -> int:
        diff = [abs(nums1[i] - nums2[i]) for i in range(len(nums1))]
        total = k1 + k2
        if total >= sum(diff):
            return 0

        def cost(line):
            return sum([d - line for d in diff if d > line ])

        l, r = 0, max(diff)
        while l < r:
            m = (l + r) // 2
            if cost(m) <= total:
                r = m
            else:
                l = m + 1

        total -= cost(l)
        res = 0
        for d in diff:
            d = min(d, l)
            if d == l and total > 0:
                d -= 1
                total -= 1
            res += d**2
        return res