partition.rb

class Partition
  include Enumerable
  attr_reader :arr, :k

  def initialize(arr, k)
    @arr = arr
    @k = k
  end

  def each
    all_partitions(@arr) { |partition| yield partition }
  end

  private

  def partition_set(set, index, ans, &block)
    if index == set.size
      yield ans.map(&:dup) if ans.size <= @k
      return
    end

    ans.each do |subset|
      subset << set[index]
      partition_set(set, index + 1, ans, &block)
      subset.pop
    end

    ans << [ set[index] ]
    partition_set(set, index + 1, ans, &block)
    ans.pop
  end

  def all_partitions(set, &block)
    partition_set(set, 0, [], &block)
  end
end