Kotlinで順列を全探索する

こんにちは、bitFlyerのAdvent Calendar7日目の記事です。
Androidエンジニアをやっています。

突然ですが皆さんは、順列を全探索するプログラムを書いたことはありますでしょうか。
例えば、[1, 2, 3]というリストが与えられたとき

1 2 3
1 3 2
2 1 3
2 3 1
3 1 2
3 2 1

を出力するようなプログラムです。

簡単そうなのに実際にやると難しかったです。
LeetCodeにも同様の問題が出されていて、有名な問題なのかもしれません。

私は以前Pythonで競プロをやっていて、Pythonにはそのためのライブラリがありました。
なので特に困らなかったのですが、どうやらKotlinにはそれがなさそうなので、作ってみました。

2パターン作りました。

まず最初のパターンは、実装は容易いけどパフォーマンスに問題がある方です。

fun main() {
    val n = 3

    val a = List(n) { it }.permutations()
    println(a)
}

fun <T> List<T>.permutations(): List<List<T>> {
    if (isEmpty()) return listOf(emptyList())

    val li = mutableListOf<List<T>>()
    for (i in indices) {
        (this - this[i]).permutations().forEach {
            li.add(it + this[i])
        }
    }
    return li
}

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

こちらのメリットは

1. アルゴリズムがわかりやすい

デメリットは

1. メモリを大量に食う（一度で全ての並びが入ったリストを返してしまうため）
2. 高速ではない（再帰を使っているから or add()が原因かなと思っています）

直感的にはわかりやすいのですが、競プロでこっちを使ってしまうとTLE( Time Limit Exceeded )になりやすいです（というか実際になってしまいました）
Iteratorのnext()的なメソッドをコールして次の値だけがポンっと出てきてくれたらメモリを食わず、かつ再帰関数を使わず高速なのですが・・・

というわけでnext()的なものを実装してみました。
ただし、下記のアルゴリズムには前提条件が二つあります。

1. 昇順にソート済みであること
2. 同じ値が存在しないこと

※若干コードが長くて複雑に見えるかもしれませんが、「辞書順で並べた時の一つ先を取得しているだけ」という認識を持って読み進めればわかりやすいかと思います。

まずはコードを貼ります

fun main(args: Array<String>) {
    val n = readLine()!!.toInt()

    val arr = Array(n) { it }
    do {
        println(arr.joinToString(" "))
    } while (arr.nextPermutation())
}

fun <T : Comparable<T>> Array<T>.nextPermutation(): Boolean {
    // AtCoderのKotlinのバージョンが古すぎてこのメソッドは生えていないので自前で定義しています
    fun <T> Array<T>.reverse(start: Int, end: Int) {
        val mid = ((start + end) / 2) - 1
        if (mid < 0) return
        var reverseIndex = end - 1
        for (index in start..mid) {
            val tmp = this[index]
            this[index] = this[reverseIndex]
            this[reverseIndex] = tmp
            reverseIndex--
        }
    }

    val n = size
    var i = n - 2
    // 末尾から、最初に大小関係が昇順になっている要素を探す
    while (i >= 0) {
        if (this[i] < this[i + 1]) break
        i--
    }
    // 要素が全て降順になっていたらreturn
    if (i < 0) return false

    // i番目の要素を最初に上回る要素を末尾から探す
    var j = n
    do {
        j--
    } while (this[i] > this[j])

    // 同じ値は想定しない
    if (this[i] == this[j]) throw IllegalArgumentException("elements are the same")

    // i番目の要素を最初に上回る要素をi番目の要素とswapする
    val tmp = this[i]
    this[i] = this[j]
    this[j] = tmp
    // 　
    this.reverse(i + 1, n)
    return true
}

3
0 1 2
0 2 1
1 0 2
1 2 0
2 0 1
2 1 0

nextPermutation()内でやっていることは、辞書順で次の値を見つけることなのですが、それをする上で下記の五つのステップを踏んでいます。

1. 末尾から、大小関係が昇順になっている2要素を探す
2. 1.で要素が見つからなかったとき、つまり全て降順になっているときは辞書順の最後にあたるのでfalseを返す
3. 1.において2要素が見つかったときは、2要素のうちの手前の要素を上回る要素を末尾から探す
4. 3.で見つけた値、と1.で見つけた2要素の手前の要素を入れ替える
5. 降順になっている範囲を昇順に戻す

実例を挙げると
1,3,2,4,0の次の値を見つけるには

1. 2と4を見つける
2. 2(1. で見つけた2要素の手前の要素）を上回る値を末尾から探すと、4が見つかる
3. 1. で昇順に並ぶ2要素が見つかったので後続の処理に続く
4. 2と4を入れ替える(1,3,4,2,0になる)
5. 降順になっている範囲を昇順に戻す
   1,3,4,0,2という辞書順の次の値が見つかります。

辞書順に全探索するというのは盲点でしたが、「どこまで探索したか」という情報を値自体が自ずと保持してくれるのでnext()を実装したいときは非常に便利だなと思いました。
ここまでで一件落着なわけですが、リスト内の値の重複が許されないのが胸につっかえています。

少し変更するだけで同じ値が入っていた場合も辞書順が出力されるようにできたのでご共有します。

変更するポイントは二つだけで
一つは等しい要素があったら例外を投げる部分を消去
もう一つは
Before

    // i番目の要素を最初に上回る要素を末尾から探す
    var j = n
    do {
        j--
    } while (this[i] > this[j])

After

    // i番目の要素を最初に上回る要素を末尾から探す
    var j = n
    do {
        j--
    } while (this[i] >= this[j])

※>を>=に変えただけです。

結論としては、等しい値をswapせずに飛ばすようにしただけです（等しい値をswapしても意味がないので）
まずここのwhile文が何をする部分か思い出してみると、「降順で並んだ2要素の手前の要素(i番目とする）と、それを上回る要素をswapするためのindexを探す部分」でした。
「i番目を上回る要素」を探していましたが、i番目と等しい要素はスキップするアルゴリズムになっています。

重複を含む順列の全列挙もこれを用いてできます。

長くなりましたが、以上になります。
Kotlinでの競プロは慣れてくると業務コードを書くスピードが上がるかも、と密かに期待しています。

最後に、弊社bitFlyerではAndroidエンジニアを募集しています。ご興味のある方はぜひご応募ください！