Knuth Algorithm U Set Partitions

gistfile1.rb

class Partitioner

  def initialize(ns, m)
    @ns = ns
    @m = m
  end

  def partitions()
    n = @ns.length
    a = [0] * (n + 1)
    (1...@m + 1).each do |j|
      a[n - @m + j] = j - 1
    end
    f(@m, n, 0, n, a)
  end

  def visit(n, a)
    ps = (0...@m).map { [] }
    for j in (0...n)
      ps[a[j + 1]] << @ns[j]
    end
    ps
  end

  def b(mu, nu, sigma, n, a)
    return Enumerator.new do |enum|
      if nu == mu + 1
        while a[nu] < mu - 1 do
          a[nu] = a[nu] + 1
        end
        a[mu] = 0
      elsif nu > mu + 1
        if (a[nu] + sigma) % 2 == 1
          for v in f(mu, nu - 1, 0, n, a)
            enum.yield v
          end
        else
          for v in b(mu, nu - 1, 0, n, a)
            enum.yield v
          end
        end
        while a[nu] < mu - 1 do
          a[nu] = a[nu] + 1
          if (a[nu] + sigma) % 2 == 1
            for v in f(mu, nu - 1, 0, n, a)
              enum.yield v
            end
          else
            for v in b(mu, nu - 1, 0, n, a)
              enum.yield v
            end
          end
        end
        if (mu + sigma) % 2 == 1
          a[nu - 1] = 0
        else
          a[mu] = 0
        end
      end
      if mu != 2
        for v in b(mu - 1, nu - 1, (mu + sigma) % 2, n, a) do
          enum.yield v
        end
      end
    end
  end

  def f(mu, nu, sigma, n, a)
    return Enumerator.new do |enum|
      if mu == 2
        enum.yield visit(n, a)
      else
        for v in f(mu - 1, nu - 1, (mu + sigma) % 2, n, a)
          enum.yield v
        end
      end
      if nu == mu + 1
        a[mu] = mu - 1
        enum.yield visit(n, a)
        while a[nu] > 0 do
          a[nu] = a[nu] - 1
          enum.yield visit(n, a)
        end
      elsif nu > mu + 1
        if (mu + sigma) % 2 == 1
          a[nu - 1] = mu - 1
        else
          a[mu] = mu - 1
        end
        if (a[nu] + sigma) % 2 == 1
          for v in b(mu, nu - 1, 0, n, a)
            enum.yield v
          end
        else
          for v in f(mu, nu - 1, 0, n, a)
            enum.yield v
          end
        end
        while a[nu] > 0 do
          a[nu] = a[nu] - 1
          if (a[nu] + sigma) % 2 == 1
            for v in b(mu, nu - 1, 0, n, a)
              enum.yield v
            end
          else
            for v in f(mu, nu - 1, 0, n, a)
              enum.yield v
            end
          end
        end
      end
    end
  end
end