Today’s problem

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

Intuition

For each point, is it easy to find out that the answer of that point is the sum of the distance between i and all the 1s.
That is: $ans[i]=\sum_{j=0}^{n-1} boxes[j] \times abs(j-i)$.
Now, imagine that there is a pointer moving left to right from 1 to n, calculating the result.
We can find that for $ans[i]$ and $ans[i+1]$, the difference would be the number of 1s from 1 to i minus the number of 1s from i+1 to n.
That is: $ans[i+1] - ans[i] = \sum_{j=0}^{i} boxes[j] - \sum_{j=i+1}^{n-1} boxes[j]$. Therefore, by calculating the number of 1s on the left hand side of i and the number of all 1s in the sequence, we can calculate all answers by $O(N)$.

Approach

First we need to calculate $ans[0]$ and the number of all 1s in the sequence by using the equation $ans[0]=\sum_{j=1}^{n} boxes[j] \times j$. Therefore, we need to iterate through the whole sequence. Next we need to calculate $ans[i+1] - ans[i] = \sum_{j=0}^{i} boxes[j] - \sum_{j=i+1}^{n-1} boxes[j]$ for every i from 1 to n-1. Because $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]$. Therefore, all we have to. do is to count the 1s in our iteration to our answer, then apply the fomular above.

Complexity

  • Time complexity: $O(N)$, N is the size of the boxes.

  • Space complexity: $O(N)$, as we need to store our answer.

Code

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

image.png