Search code examples
haskelltreefold

How to count the number of leaves of a plant with this fold function?


Color and Plant are datatypes, defined like this:

data Color = Red | Pink | White | Blue | Purple | Green | Yellow
   deriving (Show, Eq)

data Plant=
      Leaf
   |  Blossom Color
   |  Stalk Plant Plant
   deriving (Show, Eq)

The fold function is this:

fold_plant :: (x -> x -> x) -> (Color -> x) -> x -> Plant -> x
fold_plant s b l = fold_plant'
    where fold_plant' Leaf = l
          fold_plant' (Blossom c) = b c
          fold_plant' (Stalk lef rig) = s (fold_plant' lef) (fold_plant' right)

Now I'm supposed to implement a function amount which makes use of the fold function, takes a Plant as argument and returns the amount of leaves in the plant. The main problem I think I'm facing right now is that the fold function I'm supposed to use doesn't work like the usual fold functions. By that I mean like foldr and foldl, since those take 3 arguments: a function, an initial value and a list. But in my function fold_plant there's just none of that.

This is what I thought about using but I think I don't understand everything correctly.

amount :: Plant-> Integer
amount input = length . filter (Leaf ==) fold_plant Stalk Blossom Leaf input

Solution

  • The other answer jumps straight to the solution, then defends that solution as the correct one. In this answer, I'd like to go the other direction: to describe a method by which you can invent the solution yourself.

    I will assume that you could yourself have first written this function without the fold, just with basic pattern matching and recursion. Here's what you might have come up with (modulo some naming, perhaps):

    amount Leaf = 1
    amount (Blossom c) = 0
    amount (Stalk lef rig) = amount lef + amount rig
    

    Now compare with this simplified definition of the fold (which doesn't use a where-block helper function):

    fold_plant s b l Leaf = l
    fold_plant s b l (Blossom c) = b c
    fold_plant s b l (Stalk lef rig) = s (fold_plant s b l lef) (fold_plant s b l rig)
    

    My goodness those two definitions look similar! Let's hypothesize that amount = fold_plant s b l for some choices of s, b, and l. Looking at the first two equations:

    amount           Leaf = 1
    fold_plant s b l Leaf = l
    

    So we should choose l = 1. For the second two equations:

    amount           (Blossom c) = 0
    fold_plant s b l (Blossom c) = b c
    

    So we should choose b c = 0. And the last two equations:

    amount           (Stalk lef rig) = amount lef + amount rig
    fold_plant s b l (Stalk lef rig) = s (fold_plant s b l lef) (fold_plant s b l rig)
    

    Well, recall our hypothesis is that amount = fold_plant s b l, so

    s (fold_plant s b l lef) (fold_plant s b l rig) = s (amount lef) (amount rig)
    

    and the two equations are now:

    amount           (Stalk lef rig) = amount lef + amount rig
    fold_plant s b l (Stalk lef rig) = s (amount lef) (amount rig)
    

    So we should choose s x y = x + y. Writing down these equations all at once, we have arrived at our answer:

    amount = fold_plant s b l where
        l = 1
        b c = 0
        s x y = x + y