Pythonで挿入ソート時に必要なシフト（移動）回数を求めるプログラム

配列が与えられ、それに対して挿入ソートを実行することを考えてみましょう。挿入ソートでは、配列内の各要素を順番に取り出し、正しい位置まで移動（シフト）させていきます。この記事では、配列をソートし終えるまでに必要なシフトの総数を求めるプログラムを紹介します。シフトの総数は整数値となり、すでにソート済みの配列に対しては 0 を返します。

たとえば、入力が input_array = [4, 5, 3, 1, 2] の場合、出力は 8 になります。

[4, 5, 3, 1, 2] → 0 シフト

[4, 5, 3, 1, 2] → 0 シフト

[3, 4, 5, 1, 2] → 2 シフト

[1, 3, 4, 5, 2] → 3 シフト

[1, 2, 3, 4, 5] → 3 シフト

したがって、シフトの総数は 0 + 0 + 2 + 3 + 3 = 8 となります。

解法のアプローチ

この問題は、転倒数（逆順になっているペアの数）を求める問題として捉えることができます。挿入ソートで要素が移動される回数は、配列中の転倒数と一致するためです。以下の手順で効率的に計算していきます。

• length := 入力配列 input_arr のサイズ
• temp_arr := サイズ 1000001 のリストを作成し、すべて 0 で初期化（BIT＝フェニック木として機能します）
• ans := 0
• input_arr の各要素について、次の処理を行う：
  • val := 要素の値
  • val > 0 の間、次を繰り返す：
    • ans := ans + temp_arr[val]
    • val := val - (val AND -val)
  • val := 要素の値
  • val <= 1000000 の間、次を繰り返す：
    • temp_arr[val] := temp_arr[val] + 1
    • val := val + (val AND -val)
• ans := 全ペアの総数「length × (length − 1) ÷ 2」から ans を引いた値
• ans を返す

ここで登場する val & -val という演算は、Binary Indexed Tree（フェニック木）特有のテクニックで、値の最下位ビットを取り出すことで累積和の更新・参照を O(log m) で行えます。そのため、単純な二重ループによる O(n²) の計算よりもはるかに高速に転倒数を求められます。なお、この実装は要素が 1 以上 1000000 以下の正の整数であることを前提としている点に注意してください。

実装例

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

def solve(input_arr):
    length = len(input_arr)
    temp_arr = [0] * 1000001
    ans = 0
    for item in input_arr:
        val = item
        while val > 0:
            ans += temp_arr[val]
            val -= val & -val
        val = item
        while val <= 1000000:
            temp_arr[val] = temp_arr[val] + 1
            val += val & -val
    ans = length * (length - 1) // 2 - ans
    return ans

print(solve([4, 5, 3, 1, 2]))

入力

[4, 5, 3, 1, 2]

出力

8