Kotlin: 順列 / 組み合わせ / 重複順列 / 重複組み合わせ を まとめて関数にしてみた。（Sequence版）

記事「JavaScript: 順列 / 組み合わせ / 重複順列 / 重複組み合わせ を まとめて関数にしてみた。」の「ジェネレータ版」を Kotlin に移植した。¹
ロジックはそのままだが、関数名や引数名、引数の順番などを変更し、拡張関数にした。
また、引数チェックを追加した。²
（元記事筆者の @ttatsf さん、良い記事をありがとうございます。）

private fun <T> pcSequenceFactory(
    selecteds: List<T> = emptyList(),
    filter: (options: List<T>, i: Int) -> List<T>
): (options: List<T>, k: Int) -> Sequence<List<T>> =
    { options, k ->
        sequence {
            if (k == 0) {
                yield(selecteds)
                return@sequence
            }

            options.forEachIndexed { i, option ->
                pcSequenceFactory(selecteds + option, filter).let {
                    it(filter(options, i), k - 1)
                }.forEach {
                    yield(it)
                }
            }
        }
    }

/** 重複なしの順列 */
fun <T> List<T>.permutationWithoutRepetition(k: Int): Sequence<List<T>> {
    require(k in 0..size) { "引数 k は 0 以上かつ $size 以下でなければなりません。k: $k" }

    return pcSequenceFactory<T> { options, i ->
        options.take(i) + options.drop(i + 1)
    }(this, k)
}

/** 重複なしの組み合わせ */
fun <T> List<T>.combinationWithoutRepetition(k: Int): Sequence<List<T>> {
    require(k in 0..size) { "引数 k は 0 以上かつ $size 以下でなければなりません。k: $k" }

    return pcSequenceFactory<T> { options, i ->
        options.drop(i + 1)
    }(this, k)
}

/** 重複ありの順列 */
fun <T> List<T>.permutationWithRepetition(k: Int): Sequence<List<T>> {
    require(k >= 0) { "引数 k は 0 以上でなければなりません。k: $k" }

    return pcSequenceFactory<T> { options, i ->
        options
    }(this, k)
}

/** 重複ありの組み合わせ */
fun <T> List<T>.combinationWithRepetition(k: Int): Sequence<List<T>> {
    require(k >= 0) { "引数 k は 0 以上でなければなりません。k: $k" }

    return pcSequenceFactory<T> { options, i ->
        options.drop(i)
    }(this, k)
}

使用例：

fun main() {
    listOf(0, 1, 2).permutationWithoutRepetition(2)
        .also { println(it.toList()) }
    // > [[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1]]

    listOf(0, 1, 2).combinationWithoutRepetition(2)
        .also { println(it.toList()) }
    // > [[0, 1], [0, 2], [1, 2]]

    listOf(0, 1).permutationWithRepetition(3)
        .also { println(it.toList()) }
    // > [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]

    listOf(0, 1).combinationWithRepetition(3)
        .also { println(it.toList()) }
    // > [[0, 0, 0], [0, 0, 1], [0, 1, 1], [1, 1, 1]]
}

おまけ

関数名が長くて読みにくいので、重複あり／なしを引数で指定できるよう実装を追加する。

enum class Repetition {
    WITHOUT_REPETITION,
    WITH_REPETITION
}

fun <T> List<T>.permutation(k: Int, repetition: Repetition) = when (repetition) {
    Repetition.WITH_REPETITION ->
        permutationWithRepetition(k)
    Repetition.WITHOUT_REPETITION ->
        permutationWithoutRepetition(k)
}

fun <T> List<T>.combination(k: Int, repetition: Repetition) = when (repetition) {
    Repetition.WITH_REPETITION ->
        combinationWithRepetition(k)
    Repetition.WITHOUT_REPETITION ->
        combinationWithoutRepetition(k)
}

使用例；

import Repetition.WITH_REPETITION
import Repetition.WITHOUT_REPETITION

fun main() {
    listOf(0, 1, 2).permutation(2, WITHOUT_REPETITION)
        .also { println(it.toList()) }
    // > [[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1]]

    listOf(0, 1, 2).combination(2, WITHOUT_REPETITION)
        .also { println(it.toList()) }
    // > [[0, 1], [0, 2], [1, 2]]

    listOf(0, 1).permutation(3, WITH_REPETITION)
        .also { println(it.toList()) }
    // > [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]

    listOf(0, 1).combination(3, WITH_REPETITION)
        .also { println(it.toList()) }
    // > [[0, 0, 0], [0, 0, 1], [0, 1, 1], [1, 1, 1]]
}

／以上

---

1. 総当たりするようなコードを書くことがたまにあって、そのときに必要になる。以前に自分で重複ありの順列のコードを書いたことがあるが、出来が悪かった。 ↩

2. そうしないと、引数が不正であっても permutationWithoutRepetition 関数などを呼び出した時点ではなく、その関数が返した Sequence を使っていると例外がスローされることになり、原因が分かりにくい。 ↩