Python：指定されたインデックス間のスワップで配列をソートできるか判定する方法

問題概要

0から n−1 までの一意な値を含む未ソートの配列 nums があるとします。さらに、配列の要素を入れ替えてよいインデックスのペアを格納した別の配列 pairs が与えられます。スワップは何度でも実行できます。このとき、許可されたスワップのみを使って配列を昇順に並べ替えることができるかどうかを判定するのが目的です。

例えば、入力が nums = [6,1,7,3,0,5,4,2]、pairs = [(0,4),(6,0),(2,7)] の場合、出力は True になります。まず (2,7) をスワップして [6,1,2,3,0,5,4,7] とし、次に (6,0) をスワップして [4,1,2,3,0,5,6,7] とし、最後に (0,4) をスワップすれば [0,1,2,3,4,5,6,7] となり、完全にソートされた状態に到達できます。

解決のアプローチ

この問題はグラフ理論の考え方で効率よく解けます。スワップ可能なインデックス同士を辺で結んだ無向グラフを考えると、同じ連結成分に属するインデックス同士は何度でも自由に入れ替えられるため、その成分内の要素は任意の位置に配置できます。したがって、「各連結成分について、含まれるインデックスの集合と対応する値の集合が一致していれば全体をソートできる」という性質を利用します。具体的には、BFS（幅優先探索）で各連結成分を走査し、インデックスのリストと値のリストをそれぞれソートして比較します。

アルゴリズムの手順

1. N := nums のサイズ、P := pairs のペア数とする
2. v := N 個の空リストからなる隣接リストを作成する
3. visited := 新しい集合（訪問管理用）を作成する
4. i を 0 から P−1 まで繰り返す：
   • pairs[i] の第2インデックスを v[pairs[i] の第1インデックス] に追加する
   • pairs[i] の第1インデックスを v[pairs[i] の第2インデックス] に追加する
5. i を 0 から N−1 まで繰り返し、i が未訪問なら以下を実行する：
   • que := 両端キュー（deque）、arr_first := インデックス格納用リスト、arr_second := 値格納用リストを用意する
   • i を訪問済みにして que に追加する
   • que が空になるまで：
     • u := que の先頭要素を取り出す
     • arr_first に u、arr_second に nums[u] を追加する
     • v[u] の各要素 s が未訪問なら、訪問済みにして que に追加する
   • arr_first と arr_second をそれぞれソートする
   • 両者が一致しない場合は False を返す
6. すべての連結成分で一致していれば True を返す

実装例

from collections import deque

def solve(nums, pairs):
    N = len(nums)
    P = len(pairs)
    v = [[] for i in range(N)]
    visited = set()

    for i in range(P):
        v[pairs[i][0]].append(pairs[i][1])
        v[pairs[i][1]].append(pairs[i][0])

    for i in range(N):
        if i not in visited:
            que = deque()
            arr_first = []
            arr_second = []

            visited.add(i)
            que.append(i)

            while len(que) > 0:
                u = que.popleft()
                arr_first.append(u)
                arr_second.append(nums[u])

                for s in v[u]:
                    if s not in visited:
                        visited.add(s)
                        que.append(s)

            arr_first = sorted(arr_first)
            arr_second = sorted(arr_second)

            if arr_first != arr_second:
                return False
    return True

nums = [6,1,7,3,0,5,4,2]
pairs = [(0,4),(6,0),(2,7)]
print(solve(nums, pairs))

入力

[6,1,7,3,0,5,4,2], [(0,4),(6,0),(2,7)]

出力

True

計算量について

このアルゴリズムでは、グラフの構築に O(P)、すべての頂点と辺を一度ずつ訪問する BFS に O(N + P)、各連結成分のソートには全体で O(N log N) かかるため、合計の時間計算量は O(N log N + P) となります。メモリ使用量は隣接リストと訪問管理のため O(N + P) です。