徒然

思いついたら書きます

Pythonで使えるnext_permutationを実装してみた

今まで配列列挙ではitertools.permutationsに任せきりだったのですが、仕組みを知りたかったので自作してみました。

辞書順で次の組み合わせが欲しいとか、重複なしの組み合わせを列挙したいとか*1C++に慣れているのでnext_permutation準拠のコードを書きたいとかいったときにお使いください。

記事の後ろにprev_permutationもあります。

仕組み

以下の記事で詳しく説明されています。

`next_permutation` の解説と、それと同様のアルゴリズムで全列挙できるもの一覧 - ブログ名

next_permutationがイマイチよくわからなかったのでまとめてみた - Qiita

以下が概要です。

  • 配列を後ろから見て、初めて A_{i} \lt A_{i+1}となる iを探す(なければreturn False)

  • 配列を再度後ろから見て、初めて A_{i} \lt A_{j}となる jを探す

  •  A_{i} A_{j}を入れ替える

  •  A_{i+1}, A_{i+2}, ..., A_{-1}を反転する

  • return True

最悪計算量はO(n)ですが、ならし計算量はO(1)らしいです(本当?)。

実装

範囲[l, r)を指定可能です。*2

def next_permutation(a: list, l: int = 0, r: int = None) -> bool:
    # a[l,r)の次の組み合わせ
    if r is None:
        r = len(a)
    for i in range(r - 2, l - 1, -1):
        if a[i] < a[i + 1]:
            for j in range(r - 1, i, -1):
                if a[i] < a[j]:
                    a[i], a[j] = a[j], a[i]
                    p, q = i + 1, r - 1
                    while p < q:
                        a[p], a[q] = a[q], a[p]
                        p += 1
                        q -= 1
                    return True
    return False

使い方

Pythonには後方評価のループがないため少し注意が必要です。

n = 3
a = list(range(n))
while True:
    print(a)
    if not next_permutation(a, 0, n):
        break

結果

[0, 1, 2]
[0, 2, 1]
[1, 0, 2]
[1, 2, 0]
[2, 0, 1]
[2, 1, 0]

Verify

onlinejudge.u-aizu.ac.jp

使用例

atcoder.jp

最後に

このアルゴリズムを最初に思いついた人は天才かと思いました。
PythonならN = 9、PyPyならN = 10くらいまでAtCoderで使えると思います。
本当はCOBOLで実装するために調べたけど、実装するのが面倒になった。

おまけ①

辞書順手前を求めるprev_permutationも置いておきます。

def prev_permutation(a: list, l: int = 0, r: int = None) -> bool:
    # a[l,r)の前の組み合わせ
    if r is None:
        r = len(a)
    for i in range(r - 2, l - 1, -1):
        if a[i] > a[i + 1]:
            for j in range(r - 1, i, -1):
                if a[i] > a[j]:
                    a[i], a[j] = a[j], a[i]
                    p, q = i + 1, r - 1
                    while p < q:
                        a[p], a[q] = a[q], a[p]
                        p += 1
                        q -= 1
                    return True
    return False

atcoder.jp

おまけ②

誰得ですが、COBOLでも書きました。

クリックで表示

       IDENTIFICATION DIVISION.
       PROGRAM-ID. PERM.
       ENVIRONMENT DIVISION.
       INPUT-OUTPUT SECTION.
       FILE-CONTROL.
           SELECT SYSIN ASSIGN TO KEYBOARD ORGANIZATION LINE SEQUENTIAL.
       DATA DIVISION.
       FILE SECTION.
       WORKING-STORAGE SECTION.
           01 WK.
               03 N PIC 9(18).
               03 I PIC 9(18).
           01 AL.
               03 AI OCCURS 1 TO 20 TIMES DEPENDING ON N.
                   05 A PIC 9(2).
           01 PERM-WK.
               03 PERM-I PIC 9(18).
               03 PERM-J PIC 9(18).
               03 PERM-P PIC 9(18).
               03 PERM-Q PIC 9(18).
               03 PERM-TMP PIC 9(18).
               03 NN PIC 9(18).
               03 COND PIC 9(1) VALUE 1.
       PROCEDURE DIVISION.
           ACCEPT N.
           PERFORM VARYING I FROM 1 BY 1 UNTIL I > N
               MOVE I TO A(I)
           END-PERFORM.
           PERFORM UNTIL COND = 0
               DISPLAY AL
               PERFORM NEXT_PERMUTATION
           END-PERFORM.
           STOP RUN.
       NEXT_PERMUTATION SECTION.
           MOVE 0 TO COND.
           COMPUTE NN = N - 1.
           PERFORM VARYING PERM-I FROM NN BY -1 UNTIL PERM-I <= 0
               IF A(PERM-I) < A(PERM-I + 1) THEN
                   MOVE 1 TO COND
                   EXIT PERFORM
               END-IF
           END-PERFORM.
           IF COND = 1 THEN
               PERFORM VARYING PERM-J FROM N BY -1
                       UNTIL A(PERM-I) < A(PERM-J)
               END-PERFORM
               MOVE A(PERM-I) TO PERM-TMP
               MOVE A(PERM-J) TO A(PERM-I)
               MOVE PERM-TMP TO A(PERM-J)
               COMPUTE PERM-P = PERM-I + 1
               MOVE N TO PERM-Q
               PERFORM UNTIL PERM-P >= PERM-Q
                   MOVE A(PERM-P) TO PERM-TMP
                   MOVE A(PERM-Q) TO A(PERM-P)
                   MOVE PERM-TMP TO A(PERM-Q)
                   ADD 1 TO PERM-P
                   SUBTRACT 1 FROM PERM-Q
               END-PERFORM
           END-IF.
       EXIT SECTION.
       END PROGRAM PERM.

使用例

atcoder.jp

*1:more-itertools.distinct_permutationsがAtCoderで使えれば良いけど2022/1現在は使えない2023/8の言語アップデートから入りました!

*2:2022/11/8 [l, r]から変更