Description Link to heading

1769.minimum-number-of-operations-to-move-all-balls-to-each-box

Solution Link to heading

Violent solution: $\Theta(n^2)$

We could use prefix sum to reduce the time complexity.

nums[i] denotes the number of ball in first i + 1 boxes, res[i] is the minimum number of operation to ove ll balls to the i + 1th box, sum is the total number of balls. Compared to moving all balls to i th box, if we want to move all ball to the i + 1th box, the balls in 0 => i - 1 all need move an additional step, while the balls in i => n - 1 will move one step less.

So: res[i] = res[i - 1] + nums[i - 1] - (sum - nums[i - 1]);

Code Link to heading 

class Solution {
public:
    vector<int> minOperations(string boxes) {
        vector<int> nums(boxes.size(), 0);
        int sum = boxes[0] - '0';
        nums[0] = boxes[0] - '0';
        for (int i = 1; i < boxes.size(); i++) {
            if (boxes[i] == '1') {
                nums[i] = nums[i - 1] + 1;
                sum++;
            } else
                nums[i] = nums[i - 1];
        }
        vector<int> res(boxes.size(), 0);
        for (int i = 1; i < boxes.size(); i++) {
            res[0] += i * (boxes[i] - '0');
        }
        for (int i = 1; i < boxes.size(); i++) {
            res[i] = res[i - 1] + nums[i - 1] - (sum - nums[i - 1]);
        }
        return res;
    }
};