【Python】Union-Findでスワップ後のハミング距離を最小化する方法

問題概要

同じ長さを持つ2つの整数配列 src と tgt、および配列 allowedSwaps が与えられるとします。allowedSwaps[i] にはペア (ai, bi) が含まれており、これは配列 src のインデックス ai にある要素と、インデックス bi にある要素を入れ替えられることを意味します。なお、特定のペアのインデックスは何度でも、任意の順序で入れ替えて構いません。

ここで、同じ長さの2つの配列におけるハミング距離とは、要素が互いに異なる位置の個数のことです。配列 src に対して任意の回数のスワップ操作を実行した後の、src と tgt の最小ハミング距離を求めます。

例

入力が src = [2,3,4,5]、tgt = [3,2,5,6]、allowedSwaps = [[0,1],[2,3]] の場合、出力は 1 になります。まずインデックス0と1を入れ替えると src = [3,2,4,5] となり、続いてインデックス2と3を入れ替えると src = [3,2,5,4] となるためです。このとき src と tgt が異なるのはインデックス3の1箇所だけなので、ハミング距離は1になります。

解法のアプローチ

この問題は Union-Find(素集合データ構造) を使うことで効率的に解けます。スワップ可能なインデックス同士を同じグループ(連結成分)にまとめ、グループごとに src と tgt の要素の出現頻度を比較することで、一致させられない要素の数を算出します。

具体的な手順は以下の通りです。

• graph を src と同じサイズのリストとして作成し、各要素に自分自身のインデックスを格納します(初期状態では全要素が独立したグループ)。
• 関数 find() を定義します。引数は x です。
  • graph[x] != x である限り、graph[x] := graph[graph[x]](経路圧縮)を行い、x := graph[x] で更新します。
  • 最終的な x(根のインデックス)を返します。
• 関数 union() を定義します。引数は x, y です。
  • x1 := find(x)、y1 := find(y) を求め、graph[x1] := y1 として2つのグループを統合します。
• メイン処理では、以下を実行します。
  • allowedSwaps 内の各ペア (x, y) に対して union(x, y) を呼び出します。
  • groups を「値がリスト型のマップ」として用意します(デフォルトは空リスト)。
  • i を 0 から len(src)-1 までループし、i1 := find(i) を求めて groups[i1] の末尾に i を追加します。
  • ans := 0 と初期化します。
  • groups のすべての値(ids)に対して、次を繰り返します。
    • counter を空のマップとして用意します。
    • ids 内の各 idx について、counter[src[idx]] += 1、counter[tgt[idx]] -= 1 を実行します。
    • 各グループの集計後、ans += sum(abs(val) for val in counter.values()) / 2 を加算します。
• ans を返します。

実装例

理解を深めるために、以下のPythonコードを見てみましょう。

from collections import defaultdict, Counter
def solve(src, tgt, allowedSwaps):
    graph = [ n for n in range(len(src)) ]

    def find(x):
        while graph[x] != x:
            graph[x] = graph[graph[x]]
            x = graph[x]
        return x

    def union(x, y):
        x1, y1 = find(x), find(y)
        graph[x1] = y1

    for x, y in allowedSwaps:
        union(x,y)

    groups = defaultdict(list)
    for i in range(len(src)):
        i1 = find(i)
        groups[i1].append(i)

    ans = 0
    for ids in groups.values():
        counter = Counter()
        for idx in ids:
            counter[src[idx]] += 1
            counter[tgt[idx]] -= 1
        ans += sum( abs(val) for val in counter.values())/2
    return ans

src = [2,3,4,5]
tgt = [3,2,5,6]
allowedSwaps = [[0,1],[2,3]]
print(solve(src, tgt, allowedSwaps))

入力

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

出力

1

アルゴリズムのポイント

各グループ内では、src 側の要素の出現回数を +1、tgt 側の出現回数を -1 としてカウンターに記録していきます。集計後にカウンターへ残る絶対値の合計は、スワップを繰り返しても一致させられなかった要素の数を表します。これを2で割ると、そのグループがハミング距離へ寄与する不一致ペアの数になります。

Union-Findによるグループ化と経路圧縮のおかげで、全体の計算量はほぼ O((N + M)・α(N))(N は配列長、M は許可されたスワップ数、α はアッカーマンの逆関数)程度に抑えられ、大規模な入力でも高速に動作します。