How to efficiently generate all lists of length `n^2` containing `n` copies of every `x < n`?

Question

Given an integer n, how can I build the list containing all lists of length n^2 containing exactly n copies of each integer x < n? For example, for n = 2, we have:

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

This can be easily done combining permutations and nub:

f :: Int -> [[Int]]
f n = nub . permutations $ concatMap (replicate n) [0..n-1]

But that is way too inefficient. Is there any simple way to encode the efficient/direct algorithm?

Answer 1

Sure, it's not too hard. We'll start with a list of n copies of each number less than n, and repeatedly choose one to start our result with. First, a function for choosing an element from a list:

zippers :: [a] -> [([a], a, [a])]
zippers = go [] where
    go l (h:r) = (l,h,r) : go (h:l) r
    go _ [] = []

Now we'll write a function that produces all possible interleavings of some input lists. Internally we'll maintain the invariant that each [a] is non-empty; hence we'll have to establish that invariant before we start recursing. In fact, this will be wasted work in the way we intend to call this function, but for good abstraction we might as well handle all inputs correctly, right?

interleavings :: [[a]] -> [[a]]
interleavings = go . filter (not . null) where
    go [] = [[]]
    go xss = do
        (xssl, x:xs, xssr) <- zippers xss
        (x:) <$> interleavings ([xs | not (null xs)] ++ xssl ++ xssr)

And now we're basically done. All we have to do is feed in an appropriate starting list.

f :: Int -> [[Int]]
f n = interleavings (replicate n <$> [1..n])

Try it in ghci:

> f 2
[[1,1,2,2],[1,2,2,1],[1,2,1,2],[2,2,1,1],[2,1,1,2],[2,1,2,1]]