<p>You are given two positive <strong>0-indexed</strong> integer arrays <code>nums1</code> and <code>nums2</code>, both of length <code>n</code>.</p>

<p>The <strong>sum of squared difference</strong> of arrays <code>nums1</code> and <code>nums2</code> is defined as the <strong>sum</strong> of <code>(nums1[i] - nums2[i])<sup>2</sup></code> for each <code>0 &lt;= i &lt; n</code>.</p>

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

<p>Return <em>the minimum <strong>sum of squared difference</strong> after modifying array </em><code>nums1</code><em> at most </em><code>k1</code><em> times and modifying array </em><code>nums2</code><em> at most </em><code>k2</code><em> times</em>.</p>

<p><strong>Note</strong>: You are allowed to modify the array elements to become <strong>negative</strong> integers.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> nums1 = [1,2,3,4], nums2 = [2,10,20,19], k1 = 0, k2 = 0
<strong>Output:</strong> 579
<strong>Explanation:</strong> The elements in nums1 and nums2 cannot be modified because k1 = 0 and k2 = 0. 
The sum of square difference will be: (1 - 2)<sup>2 </sup>+ (2 - 10)<sup>2 </sup>+ (3 - 20)<sup>2 </sup>+ (4 - 19)<sup>2</sup>&nbsp;= 579.
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> nums1 = [1,4,10,12], nums2 = [5,8,6,9], k1 = 1, k2 = 1
<strong>Output:</strong> 43
<strong>Explanation:</strong> 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)<sup>2 </sup>+ (4 - 8)<sup>2 </sup>+ (10 - 7)<sup>2 </sup>+ (12 - 9)<sup>2</sup>&nbsp;= 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.</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>n == nums1.length == nums2.length</code></li>
	<li><code>1 &lt;= n &lt;= 10<sup>5</sup></code></li>
	<li><code>0 &lt;= nums1[i], nums2[i] &lt;= 10<sup>5</sup></code></li>
	<li><code>0 &lt;= k1, k2 &lt;= 10<sup>9</sup></code></li>
</ul>

from sortedcontainers import SortedDict
class Solution:
    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:
        diff = []
        for num1, num2 in zip(nums1, nums2):
            diff.append(abs(num1-num2))
        
        k = k1 + k2
        
        count = SortedDict(Counter(diff))
        if k - sum(diff) >= 0:
            return 0
        # print(count)
        # We get diff = [4, 4, 4, 3], count is SortedDict({3: 1, 4: 3}) k = 2
        while k > 0:
            key, value = count.popitem() # greedily decrease the count of max_diff 1 by 1 (slower). 
            # at each iteration, the key(diffKeys) of count either disappear or decrease it' value by k
            if k >= value:
                k -= value
                count[key - 1] = count.get(key - 1, 0) + value
            elif k < value:
                # for example after modifying SortedDict({3: 3, 4: 1}), k = 0
                # if count[key - 1] not exist use default value 0
                # because k < value, count[key] still exist
                count[key] = value - k
                count[key - 1] = count.get(key - 1, 0) + k
                k = 0            
        res = 0
        for key, val in count.items():
            res += key ** 2 * val     
        return res

    # or without SortedDict, but faster 
    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:
        diff = [abs(num1 - num2) for num1, num2 in zip(nums1, nums2)]
        k = k1 + k2
        if sum(diff) <= k: return 0
        diffCount = Counter(diff)
        diffKeys = sorted(diffCount.keys(), reverse = True) # [4, 3] and diff = [4, 4, 4, 3]
        i = 0
        n = len(diffCount)
        count = diffCount[diffKeys[i]]
        while i < n - 1 and (diffKeys[i] - diffKeys[i+1])*count <= k:
            k -= (diffKeys[i] - diffKeys[i+1])*count
            count += diffCount[diffKeys[i+1]] # the count of diffKeys[i+1] has been increased by count times because the number(count) of current maxdiff have been decreased to diffKeys[i+1]
            i += 1
        res = sum([diffCount[key]*key**2 for key in diffKeys[i+1:]]) # we can't use any modify operation from i+1
        def findModifiedSum(k, num, count):
            if k >= num*count: 
                return 0
            k, num = k%count, num - k//count
            return k*(num-1)**2 + (count-k)*num**2

        res += findModifiedSum(k, diffKeys[i], count)
        
        return res
35600