How to fold a tree structure in racket?

Viewed 558

I have a struct that can create a tree structure:

(struct node (value left middle right))

and another struct that defines a leaf node:

(struct emptyNode ())

How do I make a function that folds the tree into one value by summing the values together, starting with the left subtree, then middle, then right?

I was thinking of converting the tree into a list, with the order of the left subtree first, then middle, then right subtrees, and then doing a foldl on the list but am unsure on how to finish it, or if this is the right approach:

(define (treeToList tree)
  (cond
    [(node? tree) (append (node-left tree)) (treeToList (node-left tree))
                  (append (node-middle tree)) (treeToList (node-middle tree))
                  (append (node-right tree)) (treeToList (node-right tree))]
    [else ] ;do nothing, leaf node
    ))
2 Answers

So, a nice way to do this is by using Racket's pattern matching. As well as this it's also nice to use an agenda-based search so that the whole process is iterative. So the function that does the summing is going to have three arguments:

  • the current node;
  • the current agenda of nodes it needs to look at;
  • the running total;

And then it decides what to based on pattern matching these arguments, looping over the nodes and pushing & popping things from the agenda.

So given definitions of node and empty-node, and a sample tree made of them as follows:

(struct node (value left middle right))
(struct empty-node ())

(define sample-tree
  (node 1
        (node 2
              (empty-node)
              (node 3 (empty-node) (empty-node) (empty-node))
              (node -1
                    (node 1 (empty-node) (empty-node) (empty-node))
                    (empty-node)
                    (empty-node)))
        (node 10 (empty-node) (empty-node) (empty-node))
        (empty-node)))

We can write a function to sum the tree like this:

(define (sum-node-tree node-tree)
  (define/match (snt-loop thing agenda sum)
    [((empty-node) '() s)
     ;; an empty node, nothing on the agenda: we're done
     s]
    [((empty-node) (cons next more) s)
     ;; an empty node, but there is an agenda, so start on the next agenda
     ;; item
     (snt-loop next more s)]
    [((node (? number? v) l m r) a s)
     ;; a node with a value: sum the value into the total, push the middle
     ;; and right children onto the agenda, and start on the left child.
     (snt-loop l (list* m r a) (+ s v))]
    [(_ _ _)
     ;; something bad in the tree
     (error 'sum-node-tree "bogus tree")])
  (snt-loop node-tree '() 0))

then

> (sum-node-tree sample-tree)
16

You could use the data/lazytree package for this, specifically the tree-fold interface.

Using the tree example in @tfb's answer (named sample-tree), we first define a function to get the children of a node:

(define (node-children t)
  (list (node-left t)
        (node-middle t)
        (node-right t)))

Then use it to create the tree representation:

(require data/lazytree)
(define t (make-tree node-children
                     sample-tree
                     #:with-data node-value
                     #:empty-pred empty-node?))

And use tree-fold to get the result:

(tree-fold + t)
=> 16
Related