Algorithm- Sum of distances between every two nodes of a Binary Search Tree in O(n)?

Question

The question is to find out sum of distances between every two nodes of BinarySearchTree given that every parent-child pair is separated by unit distance. It is to be calculated

Answer 1

We can do this by traverse the tree two times.

First, we need three array

int []left which stored the sum of the distance of the left sub tree.

int []right which stored the sum of the distance of the right sub tree.

int []up which stored the sum of the distance of the parent tree (without the current sub tree).

So, first traversal, for each node, we calculate the left and the right distance. If the node is a leaf, simply return 0, if not, we can have this formula:

int cal(Node node){
    int left = cal(node.left);
    int right = cal(node.right);
    left[node.index] = left;
    right[node.index] = right;
    //Depend on the current node have left or right node, we add 0,1 or 2 to the final result
    int add = (node.left != null && node.right != null)? 2 : node.left != null ? 1 : node.right != null ? 1 : 0;
    return left + right + add;
}

Then for the second traversal, we need to add to each node, the total distance from his parent.

             1
            / \
           2   3
          / \
         4   5

For example, for node 1 (root), the total distance is left[1] + right[1] + 2, up[1] = 0; (we add 2 as the root has both left and right sub tree, the exact formula for it is:

int add = 0; 
if (node.left != null) 
    add++;
if(node.right != null)
    add++;

For node 2 , the total distance is left[2] + right[2] + add + up[1] + right[1] + 1 + addRight, up[2] = up[1] + right[1] + addRight. The reason there is a 1 at the end of the formula is because there is an edge from the current node to his parent, so we need to add 1. Now, I denote the additional distance for the current node is add, additional distance if there is a left subtree in parent node is addLeft and similarly addRight for right subtree.

For node 3, the total distance is up[1] + left[1] + 1 + addLeft, up[3] = up[1] + left[1] + addLeft;

For node 4, the total distance is up[2] + right[2] + 1 + addRight, up[4] = up[2] + right[2] + addRight;

So depend on the current node is a left or right node, we update the up accordingly.

The time complexity is O(n)