Flatten a tree (nested list) up to a certain depth

Question

I have a recursive algebraic data type to implement Lisp-like nested lists:

data T a = N a | L [T a] deriving (Show, Eq)

L [N 1, L [N 2, L [N 3, L [N 4], N 5], N 6], N 7] is equivalent to Lisp's (1 (2 (3 (4) 5) 6) 7). Now I want to partially flatten this nested list, i.e. down to a certain level but no more:

flattenN 0 t -- L [N 1,L [N 2,L [N 3,L [N 4],N 5],N 6],N 7]
flattenN 1 t -- L [N 1,N 2,L [N 3,L [N 4],N 5],N 6,N 7]
flattenN 2 t -- L [N 1,N 2,N 3,L [N 4],N 5,N 6,N 7]
flattenN 3 t -- L [N 1,N 2,N 3,N 4,N 5,N 6,N 7]

I implemented this function as a tree recursion, where elements are unpacked from a type constructor, concatenated, then packed back in L:

flattenN :: Int -> T a -> T a
flattenN 0 x = x
flattenN n (N x) = N x
flattenN n (L []) = L []
flattenN n (L (x:xs)) = flattenN (n-1) x +++ flattenN n (L xs)
  where N  x +++ N  y = L [N x, N y]
        N  x +++ L ys = L (N x:ys)
        L xs +++ N  y = L (xs++[N y])
        L xs +++ L ys = L (xs++ys)

To me this looks a little bit ugly and it should be simpler than that. Do you have any ideas how to implement the partial flattening of a nested list function in a different way? I would be glad to receive any answer, from minimalistic and elegant to clever and convoluted, with any feature that Haskell provides.

Daniel Wagner Daniel Wagner · Accepted Answer · 2015-02-19T16:23:09

I would write a function that flattens once, then iterate it. Like this:

values (N x) = [N x]
values (L ts) = ts

flattenOnce (L ts) = L (concatMap values ts)
flattenOnce t = t

flattenN n t = iterate flattenOnce t !! n

If you're feeling cryptic, you might also implement flattenOnce as

flattenOnce t = L (concatMap values (values t))

Flatten a tree (nested list) up to a certain depth

1 Answers