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

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

1과 0으로 이루어진 문자열이 주어진다.

특정 index로 문자열의 모든 1을 이동시켰을 때의 비용 값을 반환하는 문제다. 한 칸 움직이는 비용은 1이다.

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

0번 인덱스로 모든 1을 옮기려면 2,4,5번 인덱스의 1을 0번 인덱스로 옮겨야 하기 때문에 2+4+5 = 11이 된다

1번 인덱스로 모든 1을 옮기려면 1+3+4 = 8이 된다.

 

1. Brute Force

from typing import *


class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        # brute force
        answer = []
        box_list = list(map(int, list(boxes)))

        for i in range(len(box_list)):
            temp = 0
            for j in range(len(box_list)):
                if i != j and box_list[j] == 1:
                    temp += abs(i - j)
            answer.append(temp)

        return answer

결국 인덱스의 값이 1인 경우에 인덱스 거리의 차이를 모두 더 해주면 된다.

2중 for문을 이용해서 인덱스 거리의 절댓값을 더해주면 요구하는 값을 얻을 수 있다.

 

다만 시간이 터지려고 한다.

 

2. Prefix Sum

class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        # prefix sum
        nums = list(map(int, list(boxes)))

        right_prefix, left_prefix = [], []
        right_one_count, left_one_count = [], []
        
        for num in nums:
            left_one_count.append(left_one_count[-1] + num if left_one_count else num)
        for i in range(len(nums)):
            left_prefix.append(left_one_count[i - 1] + left_prefix[-1] if i != 0 else 0)
            
        for num in nums[::-1]:
            right_one_count.append(right_one_count[-1] + num if right_one_count else num)
        for i in range(len(nums)):
            right_prefix.append(right_one_count[i - 1] + right_prefix[-1] if i != 0 else 0)

        return list(map(lambda x, y: x + y, left_prefix, right_prefix[::-1]))

문제를 생각해보면 다음과 같다.

answer [i]의 값은 i를 기준으로 왼쪽의 모든 1을 i까지 옮기는 비용 + i를 기준으로 오른쪽의 모든 1을 i까지 옮기는 비용의 합이다.

 

아래와 같은 누적합을 이용하면 이를 해결할 수 있다.

1. 선형 탐색을 하면서 현재 인덱스에서 왼쪽의 1의 개수를 저장한다.
2. 선형 탐색을 하면서 왼쪽 누적합을 구한다.
   
   1. 어떤 index의 왼쪽 누적합은 index -1까지의 1의 개수 + index -1까지의 왼쪽 누적 합의 합이다.
3. 오른쪽은 방향만 뒤집어서 동일하게 진행한다.
4. 왼쪽, 오른쪽 누적합을 더한다.

3. 최적화

class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        left_cnt, left_cost, right_cnt, right_cost = 0, 0, 0, 0
        nums = list(map(int, list(boxes)))
        length = len(nums)
        answer = [0] * length
        for i in range(1, length):
            if nums[i - 1] == 1:
                left_cnt += 1
            left_cost += left_cnt
            answer[i] = left_cost

        for i in range(length - 2, -1, -1):
            if nums[i + 1] == 1:
                right_cnt += 1
            right_cost += right_cnt
            answer[i] += right_cost
        return answer

결국 하는 것은 2번 문제와 동일하지만 이전 1의 개수를 저장해서 선형 탐색 2번으로도 구할 수 있다.

 

생각보다 너무 시간이 오래 걸려서 실전에선 아마 못 풀었을 것 같다.