Getting all the diagonals of a matrix in Haskell

Question

The two-dimensional list looks like:

1 | 2 | 3
- - - - -
4 | 5 | 6
- - - - -
7 | 8 | 9

Or in pure haskell

[ [1,2,3], [4,5,6], [7,8,9] ]

The expected output for diagonals [ [1,2,3], [4,5,6], [7,8,9] ] is

[ [1], [4, 2], [7, 5, 3], [8, 6], [9] ]

Writing allDiagonals (to include anti-diagonals) is then trivial:

allDiagonals :: [[a]] -> [[a]]
allDiagonals xss = (diagonals xss) ++ (diagonals (rotate90 xss))

My research on this problem

Similar question here at StackOverflow

• Python this question is about the same problem in Python, but Python and Haskell are very different, so the answers to that question are not relevant to me.

• Only one This question and answer are in Haskell, but are only about the central diagonal.

Hoogle

Searching for [[a]] -> [[a]] gave me no interesting result.

Independent thinking

I think the the indexing follows a kind of count in base x where x is the number of dimensions in the matrix, look at:

1 | 2
- - -
3 | 4

The diagonals are [ [1], [3, 2], [4] ]

• 1 can be found at matrix[0][0]
• 3 can be found at matrix[1][0]
• 2 can be found at matrix[0][1]
• 1 can be found at matrix[1][1]

That is similar to counting in base 2 up to 3, that is the matrix size minus one. But this is far too vague to be translated into code.

Answer 1

Here is a recursive version, assuming that the input is always well-formed:

diagonals []       = []
diagonals ([]:xss) = xss
diagonals xss      = zipWith (++) (map ((:[]) . head) xss ++ repeat [])
                                  ([]:(diagonals (map tail xss)))

It works recursively, going from column to column. The values from one column are combined with the diagonals from the matrix reduced by one column, shifted by one row to actually get the diagonals. Hope this explanation makes sense.

For illustration:

diagonals [[1,2,3],[4,5,6],[7,8,9]]
= zipWith (++) [[1],[4],[7],[],[],...] [[],[2],[5,3],[8,6],[9]]
= [[1],[4,2],[7,5,3],[8,6],[9]]

Another version that works on the rows instead of the columns, but based on the same idea:

diagonals []       = repeat []
diagonals (xs:xss) = takeWhile (not . null) $
    zipWith (++) (map (:[]) xs ++ repeat [])
                 ([]:diagonals xss)

Compared to the specified result, the resulting diagonals are reversed. This can of course be fixed by applying map reverse.

Answer 2

Starting with universe-base-1.0.2.1, you may simply call the diagonals function:

Data.Universe.Helpers> diagonals [ [1,2,3], [4,5,6], [7,8,9] ]
[[1],[4,2],[7,5,3],[8,6],[9]]

The implementation in full looks like this:

diagonals :: [[a]] -> [[a]]
diagonals = tail . go [] where
    -- it is critical for some applications that we start producing answers
    -- before inspecting es_
    go b es_ = [h | h:_ <- b] : case es_ of
        []   -> transpose ts
        e:es -> go (e:ts) es
        where ts = [t | _:t <- b]

The key idea is that we keep two lists: a rectangular chunk we haven't started inspecting, and a pentagonal chunk (a rectangle with the top left triangle cut out!) that we have. For the pentagonal chunk, picking out the first element from each list gives us another diagonal. Then we can add a fresh row from the rectangular, un-inspected chunk to what's left after we delete that diagonal.

The implementation may look a bit unnatural, but it's intended to be quite efficient and lazy: the only thing we do to lists is destructure them into a head and tail, so this should be O(n) in the total number of elements in the matrix; and we produce elements as soon as we finish the destructuring, so it's quite lazy/friendly to garbage collection. It also works well with infinitely large matrices.

(I pushed this release just for you: the previous closest thing you could get was using diagonal, which would only give you [1,4,2,7,5,3,8,6,9] without the extra structure you want.)