Pythonで長さmの「1」グループが最後に存在するステップを効率的に求めるアルゴリズム

問題概要

1からnまでの数字の順列(パーミュテーション)を格納した配列 arr と、初期状態ですべてのビットが0に設定された長さnの2進数文字列があるとします。各ステップ i(1からnまで、2進数文字列と arr の両方でインデックスは1始まり)において、位置 arr[i] のビットが1に設定されていきます。

さらに別の値 m が与えられ、「サイズmの1のグループ」が存在する最後のステップを求める必要があります。ここで「1のグループ」とは、左右どちらの方向にも拡張できない連続した「1」の部分文字列を指します。つまり、ちょうど長さmの1のグループが存在する最後のステップを見つけ、そのようなグループが一度も現れない場合は -1 を返します。

入力例と動作の確認

例として、arr = [3,5,1,2,4]、m = 3 の場合を考えてみましょう。このとき出力は 4 になります。初期状態の2進数文字列は "00000" であり、以下のように変化していきます。

• "00100" → グループ: ["1"]
• "00101" → グループ: ["1", "1"]
• "10101" → グループ: ["1", "1", "1"]
• "11101" → グループ: ["111", "1"]
• "11111" → グループ: ["11111"]

この結果から、サイズ3のグループが存在する最後のステップは 4 であることがわかります。

解法のアプローチ

この問題を解くために、以下の手順に従います。

• n := arr のサイズとする
• num := 0(現在存在するサイズmのグループの数)
• ans := -1(答えとなるステップ番号)
• l := サイズnの配列を0で初期化(左側の連続する1の長さを記録)
• r := サイズnの配列を0で初期化(右側の連続する1の長さを記録)
• i を 0 から n-1 まで繰り返す:
  • cur := 1(新しく設定したビットを含むグループの長さ)
  • idx := arr[i] - 1
  • r[idx] が m と等しければ、num を1減らす(結合により消えるグループ)
  • l[idx] が m と等しければ、num を1減らす
  • cur := cur + l[idx] + r[idx](左右のグループと結合した新しい長さ)
  • num += (cur == m)(新しいグループがサイズmならカウント)
  • num > 0 ならば、ans := max(ans, i + 1)
  • idx - l[idx] > 0 ならば、r[idx - l[idx] - 1] := cur
  • idx + r[idx] < n - 1 ならば、l[idx + r[idx] + 1] := cur
• ans を返す

このアルゴリズムのポイントは、各ステップで隣接する既存グループの長さを l 配列と r 配列を使って O(1) で参照できることです。これにより全体の計算量は O(n) となり、毎回文字列全体を走査してグループを数え直す O(n²) の素朴なアプローチよりも大幅に高速に動作します。

実装例(Python)

def solve(arr, m):
    n = len(arr)
    num = 0
    ans = -1
    l = [0] * n
    r = [0] * n
    for i in range(n):
        cur = 1
        idx = arr[i] - 1
        if r[idx] == m:
            num -= 1
        if l[idx] == m:
            num -= 1
        cur += l[idx] + r[idx]
        num += cur == m
        if num > 0:
            ans = max(ans, i + 1)
        if idx - l[idx] > 0:
            r[idx - l[idx] - 1] = cur
        if idx + r[idx] < n - 1:
            l[idx + r[idx] + 1] = cur
    return ans

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

入力

[3,5,1,2,4], 3

出力

4