Pythonで全てのボールを各ボックスに集めるための最小操作回数を求めるプログラム

問題の概要

「boxes」というバイナリ文字列（0と1のみで構成される文字列）があるとします。boxes[i] が '0' の場合は i 番目のボックスが空であることを、'1' の場合はそのボックスにボールが 1 個入っていることを意味します。

1 回の操作では、あるボックスから隣接するボックスへボールを 1 個移動できます。操作の結果、1 つのボックスに複数のボールが入っても構いません。ここで、サイズ n の配列 answer を求めます。answer[i] は、すべてのボールを i 番目のボックスに集めるために必要な最小操作回数です。

たとえば、入力が boxes = "1101" の場合、出力は [4, 3, 4, 5] となります。

• 1 番目のボックスに集める場合：2 番目のボックスから 1 操作、最後のボックスから 3 操作が必要なため、合計 4 操作。
• 2 番目のボックスに集める場合：1 番目のボックスから 1 操作、最後のボックスから 2 操作が必要なため、合計 3 操作。
• 3 番目のボックスに集める場合：2 番目と最後のボックスからそれぞれ 1 操作、1 番目のボックスから 2 操作が必要なため、合計 4 操作。
• 最後のボックスに集める場合：1 番目のボックスから 3 操作、2 番目のボックスから 2 操作が必要なため、合計 5 操作。

解決のためのアプローチ

各ボックスごとに毎回距離を計算し直す素朴な方法では O(n²) の計算量が必要ですが、累積和の考え方を使えば O(n) で効率的に解けます。手順は以下のとおりです。

1. left := 0、right := 0、dist := 0 で初期化します。
2. i を 0 から boxes の長さ - 1 まで繰り返します。
   • boxes[i] が "1" の場合：dist := dist + i とします（インデックス 0 までの距離の合計）。さらに、i が 0 なら left を 1 増やし、そうでなければ right を 1 増やします。
3. dist を初期値とするリスト arr を作成します。
4. i を 1 から boxes の長さ - 1 まで繰り返します。
   • arr[i-1] + left - right を arr の末尾に追加します。
   • boxes[i] が "1" の場合：left を 1 増やし、right を 1 減らします。
5. arr を返します。

このアルゴリズムが機能する理由

注目位置を 1 つ右にずらすとき、左側にあるボールはすべて 1 操作ずつ遠ざかり、右側（現在位置以降）にあるボールはすべて 1 操作ずつ近づきます。したがって、前の答えに「left（左側のボール数）− right（右側のボール数）」を加算するだけで次の答えが求まります。ボールを通過するたびに left と right のカウントを更新すれば、配列全体を線形時間で処理できます。

実装例

理解を深めるために、以下の実装を見てみましょう。

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]