Monday, August 4, 2008

Two-dimensional zip

I've been doing battle with things like CSS recently (the W3C will be first against the wall when the revolution comes,) so it was a pleasure to run into a little Haskell puzzle and an elegant solution.

I was working with a matrix represented with a list of lists: [[a]]. I wanted to add the row and column indices to each element, giving: [[((Integer,Integer),a)]].

This isn't a pattern I remember seeing before, but it seemed simple enough, so I just hacked away. Here's my first try:

> index grid = zipWith indexrow [0..] grid
>     where indexrow rn rs = zipWith (\cn r -> ((rn,cn),r)) [0..] rs

Serviceable, but neither pretty nor general. The one dimensional case is so elegant; for any list, you can add indices with:

> index1d xs = zip [0..] xs

The zip function is simple and the infinite list of indices covers all cases. So it occurred to me I was looking for two things: a two-dimensional version of zip; and a grid of index tuples, semi-infinite in both directions.

Both turned out to be simple enough. Here's the grid of tuples using list comprehensions:

> indices2d = [[(r,c) | c <- [0..]] | r <- [0..]]

And the two-dimensional zip:

> zip2d = zipWith zip

And it's easy to reproduce the "With" version:

> zip2dWith f = zipWith (zipWith f)

So now my function looks like this:

> index2d grid = zip2d indices2d grid

or just:

> index2d' = zip2d indices2d

This version is clearer, simpler and easily generalizable to higher dimensions.