Search code examples
rubyarraysiterationcross-product

Iterate set covered by cross-product of ranges in ruby


I figured this answer had been asked before, so I searched, but I couldn't find anything. Granted, there are a ton of Ruby Array questions, so it might be there, just buried.

In any case, I'm trying to reduce a cross-product of ranges, returning a sum of all elements of the cross-product that meet some set of conditions. To construct a trivial example, if I have an array like this:

[0..1,0..1,0..1]

I'd like to iterate over this set:

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

and return a sum based the condition "return 1 if i[0] == 1 and i[2] == 0" (which would give 2). In my contrived example, I could do it like this:

br = 0..1

br.reduce(0){|sumx, x|
  sumx + br.reduce(0){|sumy, y|
    sumy + br.reduce(0){|sumz, z|
      sumz + (x == 1 and z == 0 ? 1 : 0)
    }
  }
}

, but in the actual application, the set of ranges might be much larger, and nesting reduces that way would get quite ugly. Is there a better way?


Solution

  • There are two orthogonals tasks, try not to to mix them up so the code remains modular.

    1. How to build a cartesian product of N arrays.

    2. How to filter the product and count.

    Use Array#product to get the cartesian product:

    xs = [0..1, 0..1, 0..1].map(&:to_a)
    xss = xs[0].product(*xs[1..-1]) # or xs.first.product(*xs.drop(1))
    #=> [[0, 0, 0], [0, 0, 1], [0, 1, 0], ..., [1, 1, 0], [1, 1, 1]]
    

    And now do the filter and couting:

    xss.count { |x, y, z| x == 1 && z == 0 }
    #=> 2
    

    This is slightly uglier that it should, that's because we'd need the classmethod Array::product instead of the method Array#product. No problem, let's add it into our extensions module and finally write:

    Array.product(0..1, 0..1, 0..1).count { |x, y, z| x == 1 && z == 0 }
    #=> 2