Search code examples
javaalgorithmprefix-sum

Intuition behind calculating the prefix and suffix sums


I am solving a LeetCode question: 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. 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. For input boxes = "001011", the output is: [11,8,5,4,3,4].

Doing it in O(n^2) is trivial. I could only solve it that way. I am trying to understand this O(n) solution, but having a hard time:

class Solution {
    public int[] minOperations(String boxes) {
        int n = boxes.length();

        int[] left = new int[n];
        int[] right = new int[n];
        int[] ans = new int[n];

        int count = boxes.charAt(0) - '0';
        for(int i = 1 ; i < n ; i++){
            left[i] = left[i - 1] + count;
            count += boxes.charAt(i) - '0';
            // System.out.println("i: "+i+" left[i]: "+left[i]+" left[i-1] : "+left[i-1]+" count: " + count);
        }

        count = boxes.charAt(n - 1) - '0';
        for(int i = n - 2 ; i >=0 ; i--){
            right[i] = right[i + 1] + count;
            count += boxes.charAt(i) - '0';
            // System.out.println("i: "+i+" right[i]: "+right[i]+" right[i+1] : "+right[i+1]+" count: " + count);
        }
        
        for(int i = 0 ; i < n ; i++) {
            ans[i] = left[i] + right[i];
        }

        return ans;
    }
}

Could someone please elaborate the logic behind:

left[i] = left[i - 1] + count;
count += boxes.charAt(i) - '0';

I understand we increment count whenever we encounter a ball, but how does left[i] = left[i - 1] + count; help us count the number of operations needed so far to move all the balls on the left to i (and vice versa in case of right)?

Thank you!


Solution

  • This comment from @dunkypie helped:

    "I finally build the intuition for the problem using DP. Here is goes. When we say calculating the number of operations for moving all the balls to the left of a box to that bax, say we are at the i th position(or box). This consists of two parts, first dp[i - 1] will give us the number of operations to move all the balls so far till (i - 1) th position and now we have all the balls till the (i - 1) th position in the (i - 1) th position(or box). Then the next part involves moving all those balls in (i - 1) th position to the i th position. Also note the cost of moving a single ball by 1 position is 1. So the recurrence relation becomes:

    dp[i] = dp[i - 1] + (1 * balls) where 1 here is the cost of moving a single ball."