Pythonで要素を入れ替えた後、一致するペアの数を最大化するプログラム

同じ長さの数値リスト A と B、および二次元リスト C が与えられます。C の各要素は [i, j] という形式で、「A[i] と A[j] は何度でも入れ替えてよい」ということを表しています。このとき、入れ替えを行った後の A[i] = B[i] となるペアの最大数を求めるのがこの問題です。

たとえば、入力が A = [5, 6, 7, 8]、B = [6, 5, 8, 7]、C = [[0, 1], [2, 3]] の場合、出力は 4 になります。A[0] と A[1]、そして A[2] と A[3] をそれぞれ入れ替えることで、すべての位置で A[i] = B[i] を満たすことができるからです。

解き方のアプローチ

この問題のポイントは、グラフの連結成分ごとに考えることです。入れ替え操作で結ばれたインデックス同士は、そのグループ内であれば任意の順序に並べ替えることができます。したがって、各連結成分について「B 側の値の出現回数」と「A 側の値」を突き合わせ、一致する分だけカウントすれば答えが求まります。

具体的には、次の手順で進めます。

• N := A のサイズ
• graph := 与えられた辺を双方向につないだグラフを作成
• ans := 0(答えとなるカウンタ)
• seen := サイズ N の訪問済みフラグのリスト(False で初期化)
• u を 0 から N-1 まで繰り返す:
  • seen[u] が False の場合:
    • queue := キューを作成し、u を挿入
    • seen[u] := True
    • queue 内の各 node について、隣接ノード nei を確認し、未訪問なら queue に追加して seen[nei] := True とする(BFS で連結成分を収集)
    • count := queue に含まれる各 i に対する B[i] の出現回数マップ
    • queue 内の各 i について、count[A[i]] が 0 以外なら count[A[i]] を 1 減らし、ans を 1 増やす
• ans を返す

それでは、理解を深めるために実際の実装を見てみましょう。

実装例

from collections import Counter

class Solution:
    def solve(self, A, B, edges):
        N = len(A)
        graph = [[] for _ in range(N)]
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
        ans = 0
        seen = [False] * N
        for u in range(N):
            if not seen[u]:
                queue = [u]
                seen[u] = True
                for node in queue:
                    for nei in graph[node]:
                        if not seen[nei]:
                            queue.append(nei)
                            seen[nei] = True
                count = Counter(B[i] for i in queue)
                for i in queue:
                    if count[A[i]]:
                        count[A[i]] -= 1
                        ans += 1
        return ans

ob = Solution()
A = [5, 6, 7, 8]
B = [6, 5, 8, 7]
C = [[0, 1], [2, 3]]
print(ob.solve(A, B, C))

入力

[5, 6, 7, 8], [6, 5, 8, 7], [[0, 1], [2, 3]]

出力

4

解説

このアルゴリズムでは、まず BFS(幅優先探索)によって、入れ替え操作でつながっているインデックスの連結成分をすべて収集します。連結成分内では自由に並べ替えができるため、成分内の B の値の分布(Counter)に対して、A の値を一つずつ照合していきます。対応する値がまだ残っていればペアとして成立し、カウントを進めます。

計算量は、頂点数を N、辺数を E とすると O(N + E) となり、非常に効率的です。連結成分ごとに独立して処理できるため、大規模な入力にも対応できます。