Friday, April 23, 2010

Darren on Ruby: Adding arrays like vectors in Ruby

In the process of writing some RSpec tests, I found it useful to be able to add arrays of numbers together as you would vectors. Of course Ruby has a Vector class, but I was intrigued as to how I might implement such a thing and what extra utility I could provide; I came up with the following:
module Enumerable
  def **(other)
    zip(other.fill(0, other.size, size - other.size)).collect {|a| a.inject {|sum,e| sum+e}}
  end
  
  def vector_add *others
    sums = nil
    if block_given?
      each do |e|
        sums = sums.nil? ? yield(e) : sums ** yield(e)
      end
    else
      sums = self
    end
    
    others.each do |other|
      sums = sums.nil? ? other : sums ** other
    end if not others.empty?
    
    return sums
  end
end
As you can see I have overridden the ** (exponentiation) operator; not sure this is a great idea, that is perhaps violating some 'behavioral integrity', but it can always be changed or replaced. Try running the above and following code (try pasting into IRB)...
puts [1,2,3] ** [4,5,6]
puts [1,2,3,4] ** [4,5,6]
puts [1,2,3].vector_add([4,5,6])
puts [1,2,3,4].vector_add {|e| [e*1, e*2, e*3, e*4]}
puts [1,2,3,4].vector_add([4,5,6]) {|e| [e*1, e*2, e*3, e*4]}
You should get something like the following:
5
7
9
5
7
9
4
5
7
9
10
20
30
40
14
25
36
40
I'm pretty pleased with that. Now, back to RSpec...