今日题目

https://leetcode.com/problems/minimum-number-of-operations-to-move-all-balls-to-each-box/description

解题思路

对于每个位置，不难发现该位置的答案等于 i 与所有 1 的距离之和。 即：$ans[i]=\sum_{j=0}^{n-1} boxes[j] \times abs(j-i)$。 现在，想象有一个指针从左到右（1 到 n）移动，逐步计算结果。 可以发现，对于 $ans[i]$ 和 $ans[i+1]$，两者之差等于 1 到 i 之间 1 的个数减去 i+1 到 n 之间 1 的个数。 即：$ans[i+1] - ans[i] = \sum_{j=0}^{i} boxes[j] - \sum_{j=i+1}^{n-1} boxes[j]$。 因此，只需计算 i 左侧 1 的个数和序列中所有 1 的总数，便可以 $O(N)$ 的复杂度求出所有答案。

算法步骤

首先利用公式 $ans[0]=\sum_{j=1}^{n} boxes[j] \times j$ 计算 $ans[0]$ 以及序列中所有 1 的总数，这需要遍历整个序列。 接下来对每个 i（从 1 到 n-1），利用 $ans[i+1] - ans[i] = \sum_{j=0}^{i} boxes[j] - \sum_{j=i+1}^{n-1} boxes[j]$ 递推计算。 由于 $ans[i+1] - ans[i] = \sum_{j=1}^{i} boxes[j] - \sum_{j=i+1}^{n} boxes[j] = 2 \times \sum_{j=0}^{i} boxes[j] - \sum_{j=0}^{n-1} boxes[j]$， 因此在遍历过程中维护左侧 1 的计数，再套用上述公式即可。

复杂度分析

• 时间复杂度：$O(N)$，N 为 boxes 的长度。

• 空间复杂度：$O(N)$，用于存储答案数组。

代码

class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        now = 0
        cnt1 = 0
        for i in range(len(boxes)):
            if boxes[i] == '1':
                now += i
                cnt1 += 1

        ans = [now] * len(boxes)

        now_cnt1 = 0
        for i in range(1, len(boxes)):
            if boxes[i-1] == '1':
                now_cnt1 += 1
            ans[i] = ans[i-1] + 2 * now_cnt1 - cnt1
        
        return ans