bit全探索 in Kotlin

最近AtCoderなどに参加していて、すべての組み合わせを生成しつつ計算するといった機会になんどか遭遇し、毎回頑張って実装していたのだけど、bit全探索という方法があるらしいと知り、調べて、Kotlinでどう書くか考えた結果次のようになった。

import java.util.BitSet

fun bitFullSearch(n: Int): List<BitSet> = (0 until (1 shl n)).map { bit ->
    BitSet(n).apply {
        repeat(n) { i ->
            set(i, bit and (1 shl i) > 0)
        }
    }
}

たとえば bitFullSearch(4) などと呼び出すと、それぞれ次のbitが立ったBitSetのリストが手に入る。

{}
{0}
{1}
{0, 1}
{2}
{0, 2}
{1, 2}
{0, 1, 2}
{3}
{0, 3}
{1, 3}
{0, 1, 3}
{2, 3}
{0, 2, 3}
{1, 2, 3}
{0, 1, 2, 3}

bitSet.get(i) でそのインデックスのビットが有効かBooleanが手に入るほか、bitSet.stream()で有効なインデックスのストリームが手に入るので大体いい感じにできる。