Pythonで解く：ちょうどk回の隣接スワップ後と最大k回のスワップ後に得られる数列の個数

1からnまでの自然数が順に並んだ配列Aがあるとします。このとき、「ちょうどk回の隣接スワップを行った後に得られる数列の数（S1）」と、「最大k回のスワップを行った後に得られる数列の数（S2）」をそれぞれ求めます。ここでいう隣接スワップとは、インデックスiとi+1にある要素を入れ替える操作のことです。

問題の概要と入出力例

例として、n = 3、k = 2 の場合を考えてみましょう。元の配列は [1, 2, 3] で、このときの出力は 3, 6 になります。

• ちょうど2回の隣接スワップ後：[1, 2, 3]、[2, 3, 1]、[3, 1, 2] の3通り → S1 = 3
• 最大2回のスワップ後：
  • 0回のスワップ後：[1, 2, 3]
  • 1回のスワップ後：[2, 1, 3]、[3, 2, 1]、[1, 3, 2]
  • 2回のスワップ後：[1, 2, 3]、[2, 3, 1]、[3, 1, 2]

重複を除いて6通りあるため、S2 = 6 となります。

解法のアプローチ（動的計画法）

この問題は動的計画法（DP）で効率よく解けます。鍵となるのは次の2つの性質です。

• S1（隣接スワップ）：隣接スワップを1回行うごとに、順列の反転数（転倒数）は必ず±1変化します。したがって、ちょうどk回の隣接スワップで到達できるのは「反転数がk以下で、かつkと同じ偶奇をもつ順列」だけです。そこで、長さnの順列のうち反転数がちょうどxであるものの個数A[x]をDPで求め、A[k%2], A[k%2+2], … を合計すればS1が得られます。
• S2（任意のスワップ）：任意の2要素を入れ替えられる場合、長さnの順列のうち最大x回のスワップで到達できる個数C[x]は、「n番目の要素を動かさない場合」と「動かす場合（相手はn-1通り）」に分けられ、漸化式 C[x] = C_n-1[x] + (n−1)·C_n-1[x−1] で計算できます。

答えは非常に大きな数になりうるため、p = 10^9+7 で割った余りを返します。手順をまとめると次の通りです。

1. p = 10^9+7 とする。
2. A := [1]、C := [1] で初期化する。
3. n を 2 から n まで順に処理する：
   B := A として A を [1] に初期化し直し、同様に D := C として C を [1] にする。
4. x を 1 から min(k+1, n(n−1)/2 + 1) までループし、A に (A[x−1] + B[x] − B[x−n]) mod p を追加していく（範囲外の添字は0とみなす）。これは反転数がちょうどxの順列の個数を求める漸化式です。
5. x を 1 から n−2 までループし、C に (D[x] + (n−1)·D[x−1]) mod p を追加する。最後に (n × D の末尾要素) mod p を追加する。
6. A[k%2] から A[k] まで2つおきに合計した値 mod p が S1、C[min(n−1, k)] が S2 となるので、これらを返す。

実装例

以下はPythonによる実装例です。

p = 10**9+7
def solve(n, k):
    A = [1]
    C = [1]
    for n in range(2,n+1):
        B = A
        A = [1]
        D = C
        C = [1]

        for x in range(1,min(k+1,n*(n-1)//2+1)):
            A.append((A[-1] + (B[x] if x<len(B) else 0) - (B[x-n] if 0<=x-n else 0)) % p )
        for x in range(1,n-1):
            C.append((D[x]+(n-1)*D[x-1]) % p)
        C.append(n*D[-1] % p)
    return sum(A[k%2:k+1:2]) % p,C[min(n-1,k)]

n = 3
k = 2
print(solve(n, k))

入力

3, 2

出力

3, 6

計算量の目安

外側のループがn回、内側のループが最大 min(k+1, n(n−1)/2) 回回るため、全体の時間計算量は O(n(k+n)) 程度です。各段階で直前の表のみを保持しているため、必要なメモリも少なく抑えられています。「反転数の分布（マホーニー数）」と「任意スワップでの到達可能性」という一見異なる2つの性質を、1つのDPフレームワークで同時に扱えるのがこの解法のポイントです。