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

https://leetcode.com/problems/minimum-number-of-operations-to-move-all-balls-to-each-box/description/

leetcode - Minimum Number of Operations to Move All Balls to Each Box

문제 유형: 그리디, 구간합, 라인 스위핑

문제 난이도: Medium

 

문제

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.

 

n개의 박스를 가진다. boxes[i]가 '0'이면 빈 박스를, '1'이면 공이 하나 있음을 나타낸다.

한 번의 연산으로 인접한 박스에 공을 하나 옮길 수 있다.

answer[i]에 i번째 박스에 모든 공을 모으는 데 필요한 최소 연산의 수를 구해서 반환하시오.

 

풀이

브루트 포스를 이용해도 풀리지만, O(n)으로 푸는 방법이 있다.

우선 한 번 순회하면서 가상의 -1번 박스에 대해 필요한 오른쪽 공들의 연산횟수와 오른쪽의 공개수를 누적한다.

 

그 뒤, 순차적으로 0~n-1번박스에 대해, 각 상황에 대해 업데이트한다.

i번박스에 대해서는, 우선 왼쪽공들에 대한 연산횟수에 왼쪽 공들의 개수만큼 더한다. i를 오른쪽으로 한 칸 밀었으니 왼쪽의 공들 각각에는 하나씩 연산이 더 필요해진다. 오른쪽공들에 대한 연산횟수에 오른쪽공들의 개수만큼 뺀다. 오른쪽 공들 입장에서는 한 번씩 필요한 연산이 줄었기 때문이다.

만약 boxes[i]가 '1'이라면, 왼쪽 공들의 개수를 1 추가하고, 오른쪽 공들의 개수를 1 줄인다. 다음 박스를 고려할 때부터는 왼쪽 박스에 있는 것이기 때문이다.

ans[i]에는 왼쪽 거리총합 + 오른쪽 거리총합을 담는다.

 

코드

C++

class Solution {
public:
    vector<int> minOperations(string boxes) {
        int n = boxes.size();

        vector<int> ans(n, 0);
        int lcount = 0, rcount = 0;
        int ldist = 0, rdist = 0;
        for(int i = 0; i < n; ++i) {
            if(boxes[i] == '1') {
                rdist += i+1;
                rcount++;
            }
        }

        for(int i = 0; i < n; ++i) {
            ldist += lcount;
            rdist -= rcount;
            if(boxes[i] == '1') {
                lcount++;
                rcount--;
            }
            ans[i] = ldist + rdist;
        }



        return ans;
    }
};