Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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

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.

like image 790
Mirzhan Irkegulov Avatar asked Mar 16 '23 23:03

Mirzhan Irkegulov


1 Answers

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))
like image 103
Daniel Wagner Avatar answered Apr 01 '23 03:04

Daniel Wagner