Rubyで与えられた配列の部分集合を列挙（＋α）

競技プログラミングの問題を解いていたら、単に場合の数を数え上げるだけでなく、各場合の状態を作成して処理する必要が時々出てきた。Rubyだと Array#combination や Array#repeated_permutation などが用意されているが、他のが欲しいときにパッと組み立てられず時間を取られてしまっている。

なので練習も兼ねて、そのとき欲しかったメソッドを作ってみる。

部分集合

総当たり的に解きたいときなど、各要素について含む/含まないパターンを網羅したいことがある。場合の数は 2 ** ary.size 。

Array#combination を利用

要素数は 0 から ary.size まで考えられるので、それぞれについて組合せを列挙すればいい。楽なうえほとんどの処理をRubyの内部実装に任せられるため速い。

each_subset.rb

def each_subset(ary, &block)
    0.upto(ary.size) { |n| ary.combination(n, &block) }
end

if $0 == __FILE__
    each_subset([*0..2]) { |subset| p subset }
end

output

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

二進表記を利用

列挙の順番 k （0始まり）を二進表記したときの 1 0 がそのまま各要素の含む/含まないに対応するようにしてみる。

素直にやるならこんな感じ。Rubyは Integer#[] で i ビット目の値を取得できる。

each_subset.rb

def each_subset(ary, &block)
    (1 << ary.size).times do |k|
        subset = ary.select.with_index { |_, i| k[i] == 1 }
        block.call(subset)
    end
end

if $0 == __FILE__
    each_subset([*0..2]) { |subset| p subset }
end

output

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

ary の先頭要素 0 は最下位ビット（＝偶奇）に対応させているため有無が毎回切り替わり、一方で末尾要素 2 は最上位ビットに対応させているため後半のみに固まって現れている。

実装は見ての通り回数固定の二重ループになっていて、時間計算量は O(2^N N) 。 2^N からすれば N は大したことないので、実用上（要素数が16程度）は問題ないはず。どちらかというと #select でひたすら評価していることのほうが気になる。

---

一応 O(2^N) も試しておく。再帰で実装し、 ary がひとつ少ない場合の処理に落とし込む。

※破壊的メソッドを多用してオブジェクトを使い回しているので、安全のため Array#dup で保護している。これを全て外せば倍以上速くなる。

each_subset.rb

def each_subset(ary, selected = [], &block)
    if ary.empty?
        block.call(selected.dup)  # 配列が同一オブジェクトだと危険すぎるので複製
        return
    end

    ary = ary.dup  # 入力配列を壊さないように複製

    elem = ary.pop
    each_subset(ary, selected, &block)  # elem を含まない
    selected.unshift(elem)
    each_subset(ary, selected, &block)  # elem を含む
    selected.shift
    # ary << elem  # 複製してあれば不要

    return  # ひとまず nil を返す
end

if $0 == __FILE__
    each_subset([*0..2]) { |subset| p subset }
end

分割

ary = [0, 1, 2, 3, 4] に対して、例えば [[0, 1], [2], [3, 4]] などを作りたい。

見方を変えると、「配列の各要素間について区切りを入れる/入れない」ということなので、場合の数は 2 ** (ary.size - 1) 。空配列のときは何を返すべきか謎なので、実装に都合良い形にする。

こちらも再帰で実装する。 ary から何個かをグループ化して抽出すれば、残ったものでまたグループを作ればいい。

each_split.rb

def each_split(ary, selected = [], &block)
    if ary.empty?
        block.call(selected.map(&:dup))  # 配列が同一オブジェクトだと危険すぎるので複製
        return
    end

    # ary == part1 + part2 を保つように操作していく
    part1 = []
    part2 = ary.dup  # 入力配列を壊さないように複製

    selected << part1
    until part2.empty?
        part1 << part2.shift  # 分割位置を右へずらす
        each_split(part2, selected, &block)
    end
    selected.pop

    # ary.replace(part1)  # 複製してあれば不要

    return  # ひとまず nil を返す
end

if $0 == __FILE__
    each_split([*0..4]) { |sp| p sp }
end

output

[[0], [1], [2], [3], [4]]
[[0], [1], [2], [3, 4]]
[[0], [1], [2, 3], [4]]
[[0], [1], [2, 3, 4]]
[[0], [1, 2], [3], [4]]
[[0], [1, 2], [3, 4]]
[[0], [1, 2, 3], [4]]
[[0], [1, 2, 3, 4]]
[[0, 1], [2], [3], [4]]
[[0, 1], [2], [3, 4]]
[[0, 1], [2, 3], [4]]
[[0, 1], [2, 3, 4]]
[[0, 1, 2], [3], [4]]
[[0, 1, 2], [3, 4]]
[[0, 1, 2, 3], [4]]
[[0, 1, 2, 3, 4]]