Home > Software design >  Haskell Map on Trees with any number of branches
Haskell Map on Trees with any number of branches

Time:01-19

I am trying to write a function mapTree :: (a -> b) -> Tree a -> Tree b for a data type

data Tree a = Empty | Node a [Tree a] deriving Show

I know how to do this if the tree only had 2 (or any set number) possible branches, but here I am kind of stuck. So far my code looks like this:

mapTree :: (a -> b) -> Tree a -> Tree b
mapTree _ Empty = Empty
mapTree f (Node a []) = Node (f a) []
mapTree f (Node a (t:ts)) = Node (f a) (mapTree f t : mapTree f ts)

But mapTree f ts is not allowed, since ts has type [Tree a]. Is there any way for me to rewrite this in a way such that I am not stuck with this error?

CodePudding user response:

Consider each list of length k as a separate case:

mapTree _ Empty = Empty
mapTree f (Node a []) = Node (f a) []
mapTree f (Node a [x]) = Node (f a) [mapTree f x]
mapTree f (Node a [x,y]) = Node (f a) [mapTree f x, mapTree f y]
mapTree f (Node a [x,y,z]) = Node (f a) [mapTree f x, mapTree f y, mapTree f z]
...

Look at the expression

[mapTree f x, mapTree f y, mapTree f z, ...]

You can simplify this using map, where the function being mapped is mapTree f:

[mapTree f x, mapTree f y, mapTree f z, ...]
    == map (mapTree f) [x, y, z, ...]

which lets you replace the infinite number of fixed-length cases with a single case

mapTree f (Node a xs) = Node (f a) (map (mapTree f) xs)
  • Related