Pythonで配列内のすべての1をグループ化するための最小スワップ回数を求める方法

問題の概要

0と1のみから構成されるバイナリ配列 data が与えられたとき、配列内のすべての 1 をどこか一箇所に連続して並べる(グループ化する)ために必要な最小スワップ回数を求めます。

例えば、配列が [1,0,1,0,1,0,0,1,1,0,1] の場合、出力は 3 になります。これは [0,0,0,0,0,1,1,1,1,1,1] のように、すべての 1 を隣接させることが可能だからです。

解法のアプローチ

この問題は「累積和(プレフィックスサム)」と「スライディングウィンドウ」を組み合わせることで効率的に解くことができます。

基本的な考え方は以下の通りです。まず配列全体に含まれる 1 の総数を one とします。最終的な状態では、長さ one の区間にすべての 1 が集まっていることになります。したがって、長さ one のウィンドウを配列上でスライドさせながら、ウィンドウ内に含まれる 1 の数が最大となる位置を探します。そのとき「ウィンドウ外に残る 1 の個数」が、必要なスワップ回数の最小値となります。

アルゴリズムの手順

1. one := 0、n := 配列 data の長さ と初期化します。
2. サイズ n の累積和配列 summ を作成し、すべて 0 で埋めた後、summ[0] := data[0] とします。
3. one := one + data[0] とします。
4. i を 1 から n − 1 まで繰り返します。
   • summ[i] := summ[i - 1] + data[i]
   • one := one + data[i]
5. ans := one とします。
6. left := 0、right := one - 1 とします。
7. right < n の間、以下を繰り返します。
   • left が 0 の場合は temp := summ[right]、それ以外の場合は temp := summ[right] - summ[left - 1] とします。
   • ans := min(ans, one - temp) とします。
   • right と left をそれぞれ 1 ずつ増やします。
8. ans を返します。

実装例(Python)

以下の実装例を見て、より深く理解しましょう。

class Solution(object):
   def minSwaps(self, data):
      one = 0
      n = len(data)
      summ=[0 for i in range(n)]
      summ[0] = data[0]
      one += data[0]
      for i in range(1,n):
         summ[i] += summ[i-1]+data[i]
         one += data[i]
      ans = one
      left = 0
      right = one-1
      while right <n:
         if left == 0:
            temp = summ[right]
         else:
            temp = summ[right] - summ[left-1]
         ans = min(ans,one-temp)
         right+=1
         left+=1
      return ans
ob = Solution()
print(ob.minSwaps([1,0,1,0,1,0,0,1,1,0,1]))

入力

[1,0,1,0,1,0,0,1,1,0,1]

出力

3

計算量について

このアルゴリズムは、累積和の構築に O(n)、スライディングウィンドウの走査にも O(n) しかかからないため、時間計算量は O(n) です。また、累積和配列を保持するため空間計算量も O(n) となります。全探索的に各位置でのスワップ数を数える非効率な方法と比べ、大規模な配列でも高速に動作する点が大きなメリットです。