Extending Ruby's Enumerable

One of the things I love most about Ruby is the built-in functional programming methods such as Enumerable#map, Enumerable#select, and Enumerable#inject. These allow you to write terse yet readable code with fewer side effects. This article looks at some ways to add more exotic functional operators to make Enumerable even more powerful. To start off, here’s an example of the canonical FizzBuzz problem using functional Ruby:

rules = { 3 => 'Fizz', 5 => 'Buzz' }
(1..100).each do |n|
  string = rules.inject('') { |s,rule| s += n % rule[0] == 0 ? rule[1] : '' }
  puts s != '' ? string : n.to_s
end

Here we use inject to build a string by checking each rule. If no string is built, (ie string == '') then it simply prints the number itself, n.to_s. Inject repeatedly applies a function to each element of a list and keeps track of the result, like so:

[a, b, c, d].inject { |result, element| result += function(element) }
=> function(a) + function(b) + function(c) + function(d)

Zip Map

Another useful functional operator is zip. This allows you to merge lists using a custom function. Here is fizz buzz using zip to merge three lists:

fizzes = (1..100).map { |n| n % 3 == 0 ? 'Fizz' : '' }
buzzes = (1..100).map { |n| n % 5 == 0 ? 'Buzz' : '' }
(1..100).zip(fizzes, buzzes).map { |n,f,b| f+b == '' ? n.to_s : f+b }

What if we could combine zip and map into one function? Well thanks to Ruby’s fantastic powers of introspection, we can. Let’s call this zip_map and add it to the enumerable module.

module Enumerable
  def zip_map *lists
    self.map.with_index do |elem, idx|
      args = [elem] + lists.each.with_object(idx).map(&:[])
      yield *args, idx
    end
  end
end

We can now rewrite fizzbuzz using the new zip_map as follows:

fizzes = (1..100).map { |n| n % 3 == 0 ? 'Fizz' : '' }
buzzes = (1..100).map { |n| n % 5 == 0 ? 'Buzz' : '' }
(1..100).zip_map(fizzes, buzzes) { |n, f, b| f+b == '' ? n.to_s : f+b }

Map at Level

What if we want to apply map to a list of lists? just using map will apply a function to each element of the main list, but what if we want to apply map to each element of each sublist? We can implement map_at_level like so

module Enumerable
  def map_at_level level = 1, &block
    if level <= 1
      self.map &block
    else
      self.map { |e| e.map_at_level level-1, &block }
    end
  end
end

This recursively calls itself until the lowest level is reached then maps &block to the list at that position.

Transpose

Sometimes, when we have a list of lists and we’d like to swap the columns and rows. For example, if we were trying to solve a Sudoku puzzle, we might have some functions that act on the rows, such as checking if the row contains every number from 1 to 9. If we had a transpose method, we could reuse those functions on the columns as well. Transpose can be implemented like so:

module Enumerable
  def transpose
    self.with_index do |i|
      yield self.with_object(i).map &:[]
    end
  end
end

This allows us to do things like:

[[1,2,3],['a','b','c']].transpose
=> [[1, "a"], [2, "b"], [3, "c"]]

Or in the sudoku example something like this:

def columns_valid? sudoku
  rows_valid? sudoku.transpose
end

Group and Blockify

But in sudoku you also need to check each 3x3 block. What if we had a method that turned a list of lists into a list of blocks? We want a function that does the following:

[[1,1,2,2],[1,1,2,2],[3,3,4,4],[3,3,4,4]].blockify(2,2) # splits into 2x2 'blocks'
=> [[[1,1],[1,1]],[[2,2],[2,2]],[[3,3],[3,3]],[[4,4],[4,4]]]

We can do this by first creating a group method which splits a list into sublists, then using transpose to apply this method to both the rows and the columns like so:

module Enumerable
  def group size # groups list into sublists of size 'size'
    self.each_slice(size).to_a
  end

  def blockify width, height = nil # groups 2D array into sub-2D-arrays of size widthxheight
    self.each_slice(height||width).map { |n| n.transpose.group(width) }.flatten(1)
  end
end

Flatten at Level

In our Sudoku problem we don’t need the blocks to be 3x3 lists of lists, each block can simply be a list of all the numbers in that block. Therefore we need a flatten_at_level method to flatten each block into a list.

module Enumerable
  def flatten_at_level level, depth = nil
    self.map_at_level(level) { |list| list.flatten depth }
  end
end

We can use flatten_at_level like so:

list = [[1,1,2,2],[1,1,2,2],[3,3,4,4],[3,3,4,4]].blockify(2,2)
=> [[[1,1],[1,1]],[[2,2],[2,2]],[[3,3],[3,3]],[[4,4],[4,4]]]
list.flatten_at_level(2)
=> [[[1, 1, 1, 1], [2, 2, 2, 2]], [[3, 3, 3, 3], [4, 4, 4, 4]]]

This means we can check our sudoku as follows

def blocks_valid? sudoku
  rows_valid? sudoku.blockify(3).flatten_at_level(2)
end

Assuming a sudoku object is a 9x9 list of lists, the full list of methods to check a sudoku is valid could look like this:

def is_valid? sudoku
  rows_valid?(sudoku) && columns_valid?(sudoku) && blocks_valid?(sudoku)
end

def rows_valid? sudoku
  sudoku.all? { |row| row.sort == (1..row.length).to_a }
end

def columns_valid? sudoku
  rows_valid? sudoku.transpose
end

def blocks_valid? sudoku
  rows_valid? sudoku.blockify(Math.sqrt(sudoku.length).to_i).flatten_at_level(2)
end

Notice this implementation is generalised to work with any square sudoku, not just 9x9. By only implementing rows_valid? once and recycling it for columns and blocks, we can easily reimplement our definition of a valid sudoku by simply changing the implementation for rows_valid?.

All code in this article can be found in the following Github Gist