Loop invariant for the sum of a binary tree

Question

I want to find a loop invariant for the following piece of code, I cannot seem to figure out any sort of relationship presented in this piece of code.

This algorithm's goal is to find the sum of all the elements in a binary tree. It stores the nodes in a stack, named s.

def TreeSum(root):
    res = 0
    s.push(root)
    while s.size > 0:
        node = s.pop()
        res = res + node.num
        if node.right != None:
            s.push(node.right)
        elif node.left != None:
            s.push(node.left)
    return res 

I have noticed that res is the sum of all the current nodes' parents as well as those to the left of its parents, but I don't know how to formulate that into a loop invariant.

Answer 1

This algorithm is supposed to compute the sum of the numbers in a binary tree. To prove its correctness, there are two important facts which need to be shown: firstly that it does compute the sum of the nodes that it visits, and secondly that it visits every node once. Therefore we will need to find at least two invariants, one for the proof of each of those facts.

In order to talk more formally about "visited" or "unvisited" nodes, let's add a couple of lines to the algorithm, to keep track of a list of "visited" nodes. These lines plainly don't change the algorithm's behaviour, because visited is only written to, never read from:

def TreeSum(root):
    res = 0
    s = Stack()
    s.push(root)

    visited = [] # added line to simplify proof

    while s.size > 0:
        node = s.pop()

        visited.append(node) # added line to simplify proof

        res = res + node.num
        if node.right != None:
            s.push(node.right)
        elif node.left != None:
            s.push(node.left)
    return res

I've also added a line which is necessary for the algorithm, to let s be a new stack, otherwise either s isn't defined or if it's a global variable then we can't guarantee it's empty when the algorithm starts.

The invariants needed are as follows:

1. res equals the sum of a.num for a in visited.
2. For all nodes a in visited, and all nodes c not in visited which are descendants of a, either c in s or there is a unique node b in s such that b is a descendant of a and an ancestor of c.
3. Either root in s or root in visited.

The third invariant is needed to conclude that visited contains every node when the loop terminates, and hence that res is the sum of every node. Without this, the other two invariants are satisfied when visited is empty and res is 0, but the third invariant allows us to reject that possibility.