How to implement `zipTree` using `foldTree`?

Question

How do you implement zipTree, which behaves like this

1. takes a merger, tree1, tree2
2. for each pair of node
   • if both possess value, merge with merger
   • otherwise turn into a Leaf

>>> :t zipTree
zipTree :: (a -> a -> b) -> Tree a -> Tree a -> Tree b
>>> zipTree (+) (Branch 1 Leaf (Branch 2 Leaf Leaf)) (Branch 3 Leaf Leaf)
(Branch 4 Leaf Leaf)

Using foldTree defined below

data Tree a = Leaf | Branch a (Tree a) (Tree a) deriving (Show, Eq)

foldTree :: b -> (a -> b -> b -> b) -> Tree a -> b
foldTree e _ Leaf     = e
foldTree e n (Branch a n1 n2) = n a (foldTree e n n1) (foldTree e n n2)

mapTree :: (a -> b) -> Tree a -> Tree b
mapTree f = foldTree Leaf (\x t1 t2 -> Branch (f x) t1 t2)

I only managed to implement the self-recusive version w/o using foldTree

zipTree :: (a -> a -> b) -> Tree a -> Tree a -> Tree b
zipTree f Leaf _ = Leaf
zipTree f _ Leaf = Leaf
zipTree f (Branch aroot a1 a2) (Branch broot b1 b2) = Branch (f aroot broot) (zipTree f a1 b1) (zipTree f a2 b2)

I have worked on the direction to return a callback with foldTree on tree1 to operate on tree2, but to no avail

Anyone knows how to use foldTree to implement the zipTree? Thanks in advance!

Answer 1

Starting with your definition:

zipTree :: (a -> a -> b) -> Tree a -> Tree a -> Tree b
zipTree f Leaf _ = Leaf
zipTree f _ Leaf = Leaf
zipTree f (Branch aroot a1 a2) (Branch broot b1 b2) = Branch (f aroot broot) (zipTree f a1 b1) (zipTree f a2 b2)

Let's separate out the case matches:

zipTree :: (a -> a -> b) -> Tree a -> Tree a -> Tree b
zipTree f Leaf _ = Leaf
zipTree f (Branch aroot a1 a2) t2 =
  case t2 of
    Leaf -> Leaf
    Branch broot b1 b2 -> Branch (f aroot broot) (zipTree f a1 b1) (zipTree f a2 b2)

For clarity, let's express zipTree as a function of two arguments, using lambdas:

zipTree :: (a -> a -> b) -> Tree a -> (Tree a -> Tree b)
zipTree f Leaf = \ _ -> Leaf
zipTree f (Branch aroot a1 a2) = \ t2 ->
  case t2 of
    Leaf -> Leaf
    Branch broot b1 b2 -> Branch (f aroot broot) ((zipTree f a1) b1) ((zipTree f a2) b2)

Does that help point the way to implementing it using foldTree? I've parenthesized things suggestively.

---

In a comment, you wrote

zipTree :: (a -> a -> b) -> Tree a -> Tree a -> Tree b
zipTree f atree btree = g btree
 where
   g = foldTree (const Leaf) callback atree
   callback aroot fWithA1 fWithA2 b =
     case b of
       Leaf -> Leaf
       (Branch broot b1 b2) -> Branch (f aroot broot) (fWithA1 b1) (fWithA2 b2)

Let's see if we can make this any more readable.

First, we can inline g at its call site:

zipTree f atree btree = foldTree (const Leaf) callback atree btree

Next, we can drop the explicit abstraction and application of atree and btree:

zipTree f = foldTree (const Leaf) callback

The next step is purely a matter of convention: main worker functions like your callback tend to be named go in Haskell code. In the context of a fold, I sometimes call the "we're done" function (here const Leaf) stop in contrast. Let's pause to put these all in place:

zipTree :: (a -> a -> b) -> Tree a -> Tree a -> Tree b
zipTree f = foldTree (const Leaf) go
  where
    go aroot fWithA1 fWithA2 b =
      case b of
        Leaf -> Leaf
        (Branch broot b1 b2) -> Branch (f aroot broot) (fWithA1 b1) (fWithA2 b2)

The last bit here is also a matter of taste. You don't use any of the other arguments of go when b is a leaf, so I think it's a bit easier to follow if you define the function in multiple lines (whatever you call them) rather than using a case expression.

zipTree :: (a -> a -> b) -> Tree a -> Tree a -> Tree b
zipTree f = foldTree (const Leaf) go
  where
    go _ _ _ Leaf = Leaf
    go aroot fWithA1 fWithA2 (Branch broot b1 b2) =
      Branch (f aroot broot) (fWithA1 b1) (fWithA2 b2)

This is no more correct and no more efficient than your version. I think it's a little easier to look at, but you should use whichever you find most understandable.

---

I just noticed that your definition has a more specific type than it needs to. This makes it less useful and also makes the compiler less able to detect mistakes in its implementation. Specifically, you require the two argument trees to have elements of the same type. Let's generalize:

zipTree :: (a -> b -> c) -> Tree a -> Tree b -> Tree c
zipTree f = foldTree (const Leaf) go
  where
    go _ _ _ Leaf = Leaf
    go aroot fWithA1 fWithA2 (Branch broot b1 b2) =
      Branch (f aroot broot) (fWithA1 b1) (fWithA2 b2)

I would personally go one step further in readability, using scoped type variables to give go a type signature.

-- This line goes at the very top of the file.
{-# language ScopedTypeVariables #-}

zipTree :: forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
zipTree f = foldTree (const Leaf) go
  where
    go :: a -> (Tree b -> Tree c) -> (Tree b -> Tree c) -> Tree b -> Tree c
    go _ _ _ Leaf = Leaf
    go aroot fWithA1 fWithA2 (Branch broot b1 b2) =
      Branch (f aroot broot) (fWithA1 b1) (fWithA2 b2)