Search code examples
haskelltree

Haskell Map on Trees with any number of branches


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?


Solution

  • 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)