Search code examples
listhaskelllist-comprehensionmonadsfunctor

List comprehension in Haskell of any order or dimension


Consider the following program.

chars = [" "] ++ ["A"] ++ ["B"] ++ (repeat "ABCD")

f :: Int -> [(Char,Int)]
f n = (,) <$> (chars !! n) <*> [1..3]

g :: Int -> [[(Char,Int)]]
g 1 = (\a     -> [a    ]) <$> (f 1)
g 2 = (\a b   -> [a,b  ]) <$> (f 1) <*> (f 2)
g 3 = (\a b c -> [a,b,c]) <$> (f 1) <*> (f 2) <*> (f 3)
-- g n = (\x1 x2 ... xn -> [x1,x2,...,xn]) <$> (f 1) <*> (f 2) <*> ... (f n)

How can we write g n generally for all n > 0, without typing out the expansion explicitly as above, ideally only using Prelude (and Control.Applicative if necessary)? Note that f n = f (n-1) for all n>3, hence it may be possible to define g recursively.

The output is like this (ignore the pretty printing):

> g 1
[ [ ( 'A' , 1 ) ] , [ ( 'A' , 2 ) ] , [ ( 'A' , 3 ) ] ]

> g 2
[ [ ( 'A' , 1 ) , ( 'B' , 1 ) ]
, [ ( 'A' , 1 ) , ( 'B' , 2 ) ]
, [ ( 'A' , 1 ) , ( 'B' , 3 ) ]
, [ ( 'A' , 2 ) , ( 'B' , 1 ) ]
, [ ( 'A' , 2 ) , ( 'B' , 2 ) ]
, [ ( 'A' , 2 ) , ( 'B' , 3 ) ]
, [ ( 'A' , 3 ) , ( 'B' , 1 ) ]
, [ ( 'A' , 3 ) , ( 'B' , 2 ) ]
, [ ( 'A' , 3 ) , ( 'B' , 3 ) ]
]

> g 3
[ [ ( 'A' , 1 ) , ( 'B' , 1 ) , ( 'A' , 1 ) ]
, [ ( 'A' , 1 ) , ( 'B' , 1 ) , ( 'A' , 2 ) ]
...
, [ ( 'A' , 3 ) , ( 'B' , 3 ) , ( 'D' , 3 ) ]
]

Solution

  • g n = traverse f [1 .. n]
    

    traverse is in the Prelude (at least for the last few years).

    As that's a bit non-obvious, here's how I got there:

    1. I noticed you were applying f to the numbers from 1 to n, so I started with map f [1 .. n].

    2. This produced a [[(Char, Int)]], which is the desired result type, but it needs to be sort of... turned sideways and multiplied. You want all the lists that come from non-deterministic selection of values in the inner lists. Non-deterministic selection is the essence of the Applicative instance for [], and it turns out that sequence on something of type [[a]] is exactly the operation "produce all the lists you get from combinatorially combining elements from the inner lists". This got me to sequence $ map f [1 .. n].

    3. But the pair of sequence and map is common enough that there's an operation that does both at once. sequence . map f === traverse f. So applying that rule simplified the result. traverse f [1 .. n].