1769. Minimum Number of Operations to Move All Balls to Each Box

You have n boxes. You are given a binary string boxes of length n, where boxes[i] is ‘0’ if the ith box is empty, and ‘1’ if it contains one ball.

In one operation, you can move one ball from a box to an adjacent box. Box i is adjacent to box j if abs(i - j) == 1. Note that after doing so, there may be more than one ball in some boxes.

Return an array answer of size n, where answer[i] is the minimum number of operations needed to move all the balls to the ith box.

Each answer[i] is calculated considering the initial state of the boxes.

class Solution {
public:
    vector<int> minOperations(string boxes) {
        vector<int> ans(boxes.length(), 0);
        int lower = boxes[0] == '0' ? 0 : 1, lowerCnt = boxes[0] == '0' ? 0 : 1;
        int upper = 0, upperCnt = 0;
        for(int i = 1; i < boxes.length(); i++) {
            if(boxes[i] == '1') {
                upper += i;
                upperCnt++;
            }
        }
        ans[0] = upper;
        for(int i = 1; i < ans.size(); i++) {
            if(boxes[i] == '1') {
                upperCnt--;
                upper--;
            }
            upper -= upperCnt;
            ans[i] = lower + upper;
            lower += lowerCnt;
            if(boxes[i] == '1') {
                lowerCnt++;
                lower++;
            }
        }

        return ans;
    }
};