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

题目

给一个只有 0 和 1 构成的字符串，0 表示该 box 没有球，1 表示该 box 有球
一次操作，可以移动一个球从一个 box，到相邻的 box
给出把所有球都挪到 index 的最小操作数， index 从 0 到 n-1

思路

prefix sum 和 suffix sum
把所有球都挪到 index 0，需要的最小步骤是，index 为 1 的所有 index 总和，即 answer0，记录后续 index 为 1 的box 个数 suffix_sum
此时将所有球都挪到 index 1，需要的最小步骤是，answer0 - suffix_cnt + prefix_sum，因为所有的 suffix_cnt 都少挪了一格，就可以少 suffix_cnt 步，对应着，要把前序的 prefix_sum 添加上
更新 suffix_sum、suffix_cnt，以及对应的 prefix_sum、prefix_cnt

代码

class Solution {
public:
    vector<int> minOperations(string boxes) {
        int suffix_sum = 0;
        int suffix_cnt = 0;
        for (int i = 1; i < boxes.size(); i++) {
            if (boxes[i] == '1') {
                suffix_cnt++;
                suffix_sum += i;
            }
        }
        vector<int> ans;
        ans.push_back(suffix_sum);
        int prefix_sum = 0;
        int prefix_cnt = 0;
        if (boxes[0] == '1') {
            prefix_sum++;
            prefix_cnt++;
        }
        for (int i = 1; i < boxes.size(); i++) {
            suffix_sum -= suffix_cnt;
            ans.push_back(suffix_sum + prefix_sum);
            if (boxes[i] == '1') {
                suffix_cnt--;
                prefix_cnt++;
            }
            prefix_sum += prefix_cnt;
        }
        return ans;
        
    }
};