Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 모든 공을 각 상자에 모으기 위한 최소 이동 횟수 구하기

이 문제에서는 boxes라는 이름의 이진 문자열이 주어집니다. 문자열에서 각 인덱스는 하나의 상자를 의미하며, boxes[i]'0'이면 i번째 상자가 비어 있다는 뜻이고, '1'이면 해당 상자에 공이 하나 들어 있다는 뜻입니다.

한 번의 연산으로 공 하나를 인접한 상자로 옮길 수 있으며, 연산 후에는 특정 상자에 여러 개의 공이 모여 있을 수도 있습니다. 우리가 구해야 할 것은 크기 n인 배열 answer로, answer[i]는 모든 공을 i번째 상자로 옮기는 데 필요한 최소 연산 횟수를 나타냅니다.

문제 예시

입력이 boxes = "1101"이라고 가정하면, 출력은 [4, 3, 4, 5]가 됩니다.

  • 첫 번째 상자로 모든 공을 옮기려면: 두 번째 상자에서 1번, 마지막 상자에서 3번 이동해야 하므로 총 4번의 연산이 필요합니다.
  • 두 번째 상자로 모든 공을 옮기려면: 첫 번째 상자에서 1번, 마지막 상자에서 2번 이동해야 하므로 총 3번의 연산이 필요합니다.
  • 세 번째 상자로 모든 공을 옮기려면: 두 번째 상자와 마지막 상자에서 각각 1번씩, 첫 번째 상자에서 2번 이동해야 하므로 총 4번의 연산이 필요합니다.
  • 마지막 상자로 모든 공을 옮기려면: 첫 번째 상자에서 3번, 두 번째 상자에서 2번 이동해야 하므로 총 5번의 연산이 필요합니다.

해결 접근 방식

단순하게 매 상자마다 다른 모든 상자와의 거리를 계산하면 O(n²)의 시간 복잡도가 발생합니다. 하지만 왼쪽과 오른쪽에 있는 공의 개수를 추적하면 O(n) 시간에 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • i번째 상자까지의 총 거리를 알고 있을 때, (i+1)번째 상자까지의 거리는 이전 거리에 (왼쪽 공 개수 − 오른쪽 공 개수)를 더한 값입니다. 한 칸 오른쪽으로 이동할 때 왼쪽의 모든 공은 거리가 1씩 늘고, 오른쪽의 모든 공은 거리가 1씩 줄어들기 때문입니다.

구체적인 단계는 다음과 같습니다.

  • left := 0, right := 0, dist := 0으로 초기화합니다.
  • 0부터 (상자 개수 − 1)까지 반복합니다.
    • boxes[i]가 "1"이면:
      • dist := dist + i
      • i가 0이면 left += 1, 그렇지 않으면 right += 1
  • dist를 첫 요소로 하는 배열 arr을 생성합니다.
  • 1부터 (상자 개수 − 1)까지 반복합니다.
    • arr[i-1] + left - right를 arr 끝에 추가합니다.
    • boxes[i]가 "1"이면 left += 1, right -= 1
  • arr을 반환합니다.

예제 코드

아래 파이썬 구현을 통해 더 잘 이해할 수 있습니다.

def solve(boxes):
    left = 0
    right = 0

    dist = 0
    for i in range(len(boxes)):
        if boxes[i] == "1":
            dist += i
            if i == 0:
                left += 1
            else:
                right += 1

    arr = [dist]
    for i in range(1, len(boxes)):
        arr.append(arr[i-1] + left - right)
        if boxes[i] == "1":
            left += 1
            right -= 1

    return arr

boxes = "1101"
print(solve(boxes))

입력

"1101"

출력

[4, 3, 4, 5]

시간 및 공간 복잡도

  • 시간 복잡도: O(n) — 문자열을 두 번 순회하므로 선형 시간에 해결됩니다.
  • 공간 복잡도: O(n) — 결과 배열을 저장하기 위한 추가 공간이 필요합니다.