Triangularizing a list in Haskell

Question

I'm interested in writing an efficient Haskell function triangularize :: [a] -> [[a]] that takes a (perhaps infinite) list and "triangularizes" it into a list of lists. For example, triangularize [1..19] should return

[[1,  3,  6,  10, 15]
,[2,  5,  9,  14]
,[4,  8,  13, 19]
,[7,  12, 18]
,[11, 17]
,[16]]

By efficient, I mean that I want it to run in O(n) time where n is the length of the list.

---

Note that this is quite easy to do in a language like Python, because appending to the end of a list (array) is a constant time operation. A very imperative Python function which accomplishes this is:

def triangularize(elements):
    row_index = 0
    column_index = 0
    diagonal_array = []
    for a in elements:
        if row_index == len(diagonal_array):
            diagonal_array.append([a])
        else:
            diagonal_array[row_index].append(a)
        if row_index == 0:
            (row_index, column_index) = (column_index + 1, 0)
        else:
            row_index -= 1
            column_index += 1
    return diagonal_array

---

This came up because I have been using Haskell to write some "tabl" sequences in the On-Line Encyclopedia of Integer Sequences (OEIS), and I want to be able to transform an ordinary (1-dimensional) sequence into a (2-dimensional) sequence of sequences in exactly this way.

Perhaps there's some clever (or not-so-clever) way to foldr over the input list, but I haven't been able to sort it out.

Answer 1

Make increasing size chunks:

chunks :: [a] -> [[a]]
chunks = go 0 where
    go n [] = []
    go n as = b : go (n+1) e where (b,e) = splitAt n as

Then just transpose twice:

diagonalize :: [a] -> [[a]]
diagonalize = transpose . transpose . chunks

Try it in ghci:

> diagonalize [1..19]
[[1,3,6,10,15],[2,5,9,14],[4,8,13,19],[7,12,18],[11,17],[16]]