[leetcode/lintcode 题解] 放小球问题

n 个桶中小球的个数已知,可以操作 k 次(每次从桶中取出一个球,或者添加一个球),每个桶有规定的最大容量 W[i]。求操作后两相邻桶之间的最大差值的平方的最小值。

在线评测地址: https://www.lintcode.com/problem/put-small-balls/?utm_source=sc-v2ex-fks

样例 1:

输入: 
5
6
[1,2,3,4,5]
[15,15,15,15,15]
输出: 
0

说明:

共有 5 个桶，最多操作 6 次，桶内的小球分别是 1,2,3,4,5，桶的最大容量分别是 15,15,15,15,15 。

[题解] 二分最大差值 target，判断最大差值为 target 时的最小变动球数是否小于等于最多操作次数 k 。如果最小变动球数小于等于 k，那么说明可行，继续搜寻更小的差值;否则搜寻更大的差值。

计算最大差值为 target 时的最小变动球数：对于数组 A[]的每一个桶内小球数目，我们都需要进行[1,W[i]]范围内的可能状态的转移。令 dp[i][j]表示 A[i]=j 时，A[i]与 A[i-1]差值不大于 target 所需要的最小变动球数。当 A[i]=j 时，可行的 A[i-1]的范围为[max(0, j - target)，min(W[i - 1], j + target)]。当 A[i]=j，k 在[max(1, j - target)，min(100, j + target)]范围内时，我们可以写出以下式子：

临界值： dp[0][j] = abs(j - A[0])

状态转移方程： dp[i][j] = min(dp[i][j], dp[i - 1][k] + abs(j - A[i]))

最后在所有最后一位的可能解 dp[n-1][?]中的最小值就是我们求的最小变动球数。

class Solution:
    """
    @param n: the number of buckets
    @param k: the maximum times of operations
    @param A: the number of balls in each bucket
    @param W: the maximum capacity of each bucket
    @return: return the minimum square value of the maximum difference
    """
    # 计算当相邻差值最大为 target 时的最小变动球数
    def min_cost(self, n, A, W, target):
        # dp[i][j]表示元素 A[i]=j 时，A[i]与 A[i-1]差值不大于 target 所需要的最小变动球数
        # 初始化为极大值
        dp = [[1000000000] * 101 for _ in range(n)]
        for i in range(n):
            for j in range(W[i] + 1):
                if i == 0:
                    # 临界值：第一个元素 A[0]调整为 j 的变动球数
                    dp[0][j] = abs(j - A[0])
                else:
                    # left 为 A[i]=j 时，A[i-1]与 A[i]差值不大于 target 的 A[i-1]最小值
                    # right 为 A[i]=j 时，A[i-1]与 A[i]差值不大于 target 的 A[i-1]最大值
                    left = max(0, j - target)
                    right = min(W[i - 1], j + target)
                    for k in range(left, right + 1):
                        # 当 A[i-1]=k 时，答案为 A[i-1]=k 的代价 dp[i-1][k]，加上 A[i]=j 的调整代价 abs(j-A[i])
                        dp[i][j] = min(dp[i][j], dp[i - 1][k] + abs(j - A[i]))
        mincost = 1000000000
        for i in dp[n - 1]:
            mincost = min(mincost, i)
        return mincost
        
    def getAns(self, n, k, A, W):
        ans = 1000000000
        left, right = 0, 100
        # 二分最大差值，用 minCost 计算最小变动球数
        # 如果最小变动球数小于等于 k，那么说明可行，继续搜寻更小的差值
        # 否则搜寻更大的差值
        while left + 1 < right:
            mid = left + (right - left) // 2
            if self.min_cost(n, A, W, mid) <= k:
                ans = min(ans, mid)
                right = mid
            else:
                left = mid
        if self.min_cost(n, A, W, left) <= k:
            ans = min(ans, left)
        if self.min_cost(n, A, W, right) <= k:
            ans = min(ans, right)
        return ans*ans