長さ N の配列があり、それを M 個の部分集合に均等に分割したい。

　制約を簡単に定義するとこうなる。

TL;DR

active_support/core_ext/array/grouping の Array#in_groups がそのものズバリである。

誤ったロジック

　3年前の自分が実装したロジックをざっくり書くと、それはこんなものであった。仮に N = 10, M = 3 とおいている。

def generate_array(size)
  i = 0
  Array.new(size) { i += 1 }
end

def solve(arr, m)
  group_size = (arr.size + m - 1) / m
  arr.each_slice(group_size).to_a
end

N = 10
M = 3

arr = generate_array(N)
solve(arr, M)
# => => [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10]]

　これだと一見 M 個の部分集合に分割できたようにはみえる。しかし各部分集合に属する要素数は [4, 4, 2] となっている。 [4, 3, 3] と均等に分割したいところである。

　それだけではない。3年越しに見つけた問題点は次である。

N = 0 のとき

　Array#each_slice に 0 を渡すと ArgumentError (invalid slice size) が発生してしまう。

N / group_size != M のとき

　例えば N = 6, M = 4 のときを考えてみよう。

N = 6
M = 4

arr = generate_array(N)
solve(arr, M)
# => [[1, 2], [3, 4], [5, 6]]

　4個の部分集合に分割するはずが、3つにしか分割できていない。group_size を切り上げで計算してしまっているため、部分集合あたりの個数を正しく計算できていないのである。

def solve(arr, m)
  group_size = (arr.size + m - 1) / m
  arr.each_slice(group_size).to_a
end

　要するに、誤ったロジックであるわけである。

　配列の操作であるから、標準 API に便利メソッドがありそうな気がする。で、 ruby-doc.org で Array から Enumerable から、それらしいメソッドを探してひたすら読んでいたが、代替できそうなものは見当たらない。「自前で実装するしかないのか…」と思い始めたとき、 active_support/core_ext にそのものズバリな API を見つけた。

Array#in_groups

　手短に改善後のコードを示すと、こうなる。

require 'active_support/core_ext/array/grouping'

def generate_array(size)
  i = 0
  Array.new(size) { i += 1 }
end

def solve(arr, m)
  arr.in_groups(m, false)
end

　N = 0 で空集合に対して実行しても ArgumentError を吐かずに、M個の部分集合を作ってくれる。

N = 0
M = 4

arr = generate_array(N)
solve(arr, M)
# => [[], [], [], []]

　また先ほどうまく機能しなかった N = 6, M = 4 のケースでも、部分集合の個数はあっているし、きちんと均等に分割されてもいる。

N = 6
M = 4

arr = generate_array(N)
solve(arr, M)
# => [[1, 2], [3, 4], [5], [6]]

　ここで終わってしまってはそれまでだから、せっかくなので Array#in_groups の実装も覗いてみよう。

あるべきアルゴリズム

https://github.com/rails/rails/blob/v6.0.3.4/activesupport/lib/active_support/core_ext/array/grouping.rb#L62-L86

def in_groups(number, fill_with = nil)
  # size.div number gives minor group size;
  # size % number gives how many objects need extra accommodation;
  # each group hold either division or division + 1 items.
  division = size.div number
  modulo = size % number

  # create a new array avoiding dup
  groups = []
  start = 0

  number.times do |index|
    length = division + (modulo > 0 && modulo > index ? 1 : 0)
    groups << last_group = slice(start, length)
    last_group << fill_with if fill_with != false &&
      modulo > 0 && length == division
    start += length
  end

  if block_given?
    groups.each { |g| yield(g) }
  else
    groups
  end
end

　コメントに書かれていることがすべてではあるが、言われてみると「なるほど！」と単純に納得できるだけのシンプルなアルゴリズムである。

  # size.div number gives minor group size;
  # size % number gives how many objects need extra accommodation;
  # each group hold either division or division + 1 items.
  division = size.div number
  modulo = size % number

　配列の長さを分割したい個数で除算して、ひとつの部分集合あたりの要素数を division に定義している。 modulo は剰余であり、この個数にあたる部分集合については、 division + 1 個の要素を持つことになる。

　ところで Integer#/ ではなく Integer#div を使っている。Integer#/ は 引数のクラスに応じて返り値のクラスが決定されるのに対して、 Integer#div は引数のクラスに関わらず Integer が返り値となる。つまりはこういうことである。

4 / 2.0
# => 2.0

4.div 2.0
# => 2

　Integer#div は ActiveSupport の拡張ではなく、 Ruby 標準の API である。これも知らなかった。

　本題に戻ろう。続いてのコードはこうである。

  # create a new array avoiding dup
  groups = []
  start = 0

　分割後の部分集合を格納するための配列と添字を初期化している。 #dup を忌避するとあるが、この後の処理を読めば #dup がないのは妥当であるように思われ、コメントの意図はよくわからない。

  number.times do |index|
    length = division + (modulo > 0 && modulo > index ? 1 : 0)
    groups << last_group = slice(start, length)
    last_group << fill_with if fill_with != false &&
      modulo > 0 && length == division
    start += length
  end

　length で次の部分集合の要素数を決定している。 modulo 個の部分集合は division + 1 個の要素を持つという冒頭のコメントの実装である。 Array#slice で length 個分の部分集合を添字ベースで切り出したものを、変数 last_group として定義し、ただちに groups に代入している。次行では、各部分集合の要素数を揃えるために、 fill_with を last_group に破壊的に追加している。あらかじめ last_group にこの処理を施しておくのではなく、 groups << last_group を実行した後で last_group << fill_with としているのが面白い。 if fill_with != false とあるように、 #in_groups の第二引数に false を渡しておけば、要素数は埋められない。説明は省略したが、上で書いたコードについても fill_with = false としてある。最後に start += length ですでに分割した要素数だけ添字をずらして、イテレートブロックは終了である。

  if block_given?
    groups.each { |g| yield(g) }
  else
    groups
  end

　最後にブロックを与えられたときの動作が定義されているが、これは触れるまでもないだろう。

　計算量は部分集合の個数を M として O(M) である。冒頭で制約は 1 <= M <= 100 と与えたから、配列がいくら長くてもこのアルゴリズム単体としては十分問題なく動作できる。

　と、 Array#in_groups のアルゴリズムをみて思ったことを徒然なるままに書いた。結論として、表題のアルゴリズムを Ruby で実装するのであれば ActiveSupport のサポートを借りれば十分であるし、別言語で実装するのであっても ActiveSupport の実装を参考にすれば十分満足のいくアルゴリズムができるだろう。