Minimum Number of Operations to Move All Balls to Each Box(Array,String,Prefix Sum)

https://leetcode.com/problems/minimum-number-of-operations-to-move-all-balls-to-each-box/description/?envType=daily-question&envId=2025-01-06

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.

 

Example 1:

Input: boxes = "110"
Output: [1,1,3]
Explanation: The answer for each box is as follows:
1) First box: you will have to move one ball from the second box to the first box in one operation.
2) Second box: you will have to move one ball from the first box to the second box in one operation.
3) Third box: you will have to move one ball from the first box to the third box in two operations, and move one ball from the second box to the third box in one operation.

Example 2:

Input: boxes = "001011"
Output: [11,8,5,4,3,4]

 

Constraints:

• n == boxes.length
• 1 <= n <= 2000
• boxes[i] is either '0' or '1'.

문제 이해는 되었다. 그닥 어려워 보이진 않는다.

 

1. Python

 

time complexity 측면에서 별로 좋지 않은,, for문 2중첩으로 문제를 풀긴 했는데 쪼매 성과는 짜친다.

 

python 특유의 method를 이용해서 나름 쉽게 풀어냈다.

 

class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        n = len(boxes)
        ans = [0]*n
        for i in range(n):
            for x,y in enumerate(boxes):
                if y == '1':
                    ans[i] += abs(x-i)
        return ans

 

이 역시 prefix를 이용하면 아래와 같이 풀 수 있는데,

시간절약 측면에서 대단히 효율적이다. 다만 요정도까지 바로 풀이법으로 떠오르진 않았다 ㅎ..

 

class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        answer = []
        pref = p = s = 0
        for i, el in enumerate(boxes):
            if el == '1':
                pref += i
                p += 1
        for el in boxes:
            answer.append(pref)
            if el == '1':
                p -= 1
                s += 1
            pref = pref - p + s
        return answer

 

prefix를 이용해서 C/C++ 코드도 짜보기

 

 

2. C

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
int* minOperations(char* boxes, int* returnSize) {
    int n = strlen(boxes);
    int* ans = (int*)malloc(n*sizeof(int));

    int sum = 0, pre = 0, post = 0;

    for(int i=0;i<n;i++){
        if(boxes[i]=='1'){
            sum += i;
            post += 1;
        }
    }

    for(int j=0;j<n;j++){
        ans[j] = sum;
        if(boxes[j]=='1'){
            pre += 1;
            post -= 1;
        }
        sum = sum - post + pre ;
    }
    *returnSize = n;
    return ans;
}

 

 

3. C++

class Solution {
public:
    vector<int> minOperations(string boxes) {
        int n = boxes.size();
        vector<int> ans;
        int sum = 0, pre = 0, post = 0;
        
        for(int i=0;i<n;i++){
            if(boxes[i]=='1'){
                sum += i;
                post += 1;
            }
        }

        for(int j=0;j<n;j++){
            ans.push_back(sum);
            if(boxes[j]=='1'){
                post -= 1;
                pre += 1;
            }
            sum = sum - post + pre;
        }
        return ans;
    }
};