🐥

Juliaで競技プログラミング(bit全探索をあえてiteratorで書く)

2021/01/31に公開

bit全探索で手間取ってatCoderのC問題を落とすことが多かったのでbit全探索の問題をいくつか解いて慣らした．その過程で「iteratorで書いたらめちゃめちゃ楽なのでは」という妄想を抱いたため，実装してみた．結論から言うと~~大してメリットはない~~iteratorで書くことで見た目がすっきりするbit探索がかけた．ということで以下ではiterate関数の勉強も兼ねてあえてbit全探索をiteratorで書いてみる，あえてね．

bit全探索ってなんだっけ

有限集合 $S$ の冪集合 $\mathfrak{P}(S)$ を全て調べるときに使われるテクニックです．例えば3元集合 $\{a,b,c\}$ の冪集合は $\{\{\}(=\emptyset),\{a\},\{b\},\{a,b\},\{c\},\{a,c\},\{b,c\},\{a,b,c\}\}$ となります．この3元集合の冪集合の要素数は8個だが，実は一般にn元集合の冪集合の要素数は $2^n$ となります．賢い人ならお気づきのとおり，集合の元として1個目の元がある/ない，2個目の元がある/ない，・・・というように2通り $\times$ 2通り $\times$ ・・・という形で計算されるのです．いかにもbit演算と相性が良さそうですねえ．

実際のbit全探索のコードはこんな感じになります．

#例：N元集合{1,...,N}
N = 3
array = [n for n in 1:N]
for bit in 0:1<<N-1
    subset = Int[]
    for n in 1:N
	if bit>>(n-1)&1 == 1
	    push!(subset, n)
        end
    end
    #do something
end

自分と同じ競プロ初心者であればbit演算は本当に謎に感じるかもしれません．プログラムをいっぱい書いたり，手で計算して確認したりするうちに慣れてくるとは思いますが，本当に無理な場合は単に集合に対する演算を表す記号だと思ってもいいのかもしれません．(それだとかなり苦しくなる可能性もあるけど・・・)そういった解説だとこの記事がわかりやすかったです．

iteratorってなんだっけ

Juliaにおけるiteratorというのはざっくり

for i in iter
    #do something
end

みたいにfor文を簡潔に書けるものです．(少なくとも自分自身はそんな感じのゆるふわ理解です) ここではiterがiteratorに相当します．

iteratorとして満たすべき性質は結構シンプルです．それは型に対して適切にiterate関数が定義されていることです．適切なiterate関数とは，次の二つが定義されている必要があります．

iterate(iter)：最初の要素と最初の添え字の次の添え字の組(elem, index)または(返す要素が無い場合に)nothingを返す関数．
iterate(iter, index)：現在の要素と次の添え字の組(elem, next(index))または(現在の要素が存在しない場合に)nothingを返す関数．(nextというのはお気持ちで，通常はindex+1みたいなものを想定しています)

※indexを添え字と言っていますが，もっと抽象的にiterの内部状態を表す順序集合の元であれば何でも良さそうです．(実際に書いたことが無いので自信がない)

これだけだとわかりにくいので，シンプルかつ次に活かせる例として1から $2^N$ までの整数を要素にもつPower{N}型を作ってみましょう．(なんかついてる{N}はパラメトリック型というやつです．ここでは解説はしませんが便利です．)

struct Power{N}
    Power(N::Integer) = new{N}()
end
function Main.iterate(p::Power{N}) where N
    N<1 && return nothing
    return (1, 2)
end
function Main.iterate(p::Power{N}, index) where N
    2^N<index && return nothing
    return (index, index+1)
end

for p in Power(3)
    println(p)
end
# Out
# 1
# 2
# 3
# 4
# 5
# 6
# 7
# 8

このように一見全然要素と言えるようなものを内部に持っていなくても，適切にiterate関数さえ定義してしまえばよしなにfor文を回してくれます．すばらしい！

iteratorでbit全探索を書く

bit全探索とiteratorの解説を行ったので早速bit全探索のコードをiteratorで書きなおしてみよう．

struct IntBitSet{N}
    IntBitSet(N::Integer) = new{N}()
end
struct IntBitSubSet{N}
    set::Int
    IntBitSubSet(bs::IntBitSet{T}, set::Int) where T = new{T}(set)
end 
import Main: iterate, ∈
function iterate(bs::IntBitSet{T}) where T
    T≡0 && return nothing
    obj = IntBitSubSet(bs, 0)
    return (obj, 1)
end
function iterate(bs::IntBitSet{N}, i) where N
    1<<N-1<i && return nothing
    obj = IntBitSubSet(bs, i)
    return (obj, i+1)
end
function ∈(n::Integer, be::IntBitSubSet{N}) where N
    return be.set>>(n-1)&1==1
end

IntBitSet{N}はN元集合 $\{1,...,N\}$ の冪集合の元IntBitSubSet{N}を列挙してくれる型です．最後の∈は部分集合IntBitSubSet{N}にある要素が入っているかどうかを調べる関数です．この型を使うとbit全探索はとてもシンプルに書けます．(~~型のコードがそもそもbit全探索のコードよりカロリー高いけど~~)

N = 3 #N元集合[1,...,N]
for s in IntBitSet(N)
    subset = Int[]
    for n in 1:N
        if n in s
	    push!(subset, n)
	end
    end
    # do something
end

bit操作が完全に隠蔽されて見かけ上はとても読みやすくなりました．(inと∈は同じ表現になります)

bit全探索ってなんだっけ

iteratorってなんだっけ

iteratorでbit全探索を書く

Discussion