Vertical sum of nodes in BST

Tags: , ,

Vertical sum of nodes in BST

Given a binary search tree, sum all the nodes which are at on the same vertical line. This problem is know as vertical sum problem. For example, given a binary tree as shown below
vertical sum in binary tree

the output will be:
Sum of vertical line 1 = 2
Sum of vertical line 2 = 3
Sum of vertical line 3 = 18
Sum of vertical line 4 = 8
Sum of vertical line 5 = 9

Vertical sum in binary tree : algorithm

To add the node to a particular line we have to reach to that node. Question is what kind of traversal will be best suited for this problem? Imagine we have a binary which has only one level, root. It will be easy to say that the sum of the only vertical line is the root itself. What if there are two levels? Then we have three lines, first contains the left child of the root, second contains root itself and third the right child.

As increase the levels, the number of lines also increase and we can easily know that by moving the counter, increment it when we move to right and decrement when moving to left. So, it requires level order traversal of the binary tree.

package AlgorithmsAndMe;
import java.util.*;

public class VerticalSum {
    public Map<Integer,Integer> getVerticalSum(TreeNode root){

        class VerticalInfo {
            TreeNode node;
            int line;

            public VerticalInfo(TreeNode node, int line){
                this.node = node;
                this.line = line;
            }
        };

        Queue<VerticalInfo> queue = new LinkedList<>();
        Map<Integer, Integer> verticalSum = new HashMap<>();

        if(root == null) return verticalSum;

        queue.add(new VerticalInfo(root, 0));

        while (!queue.isEmpty()){
            VerticalInfo current = queue.poll();
            if(!verticalSum.containsKey(current.line)){
                verticalSum.put(current.line,0);
            }

            /* add to the corresponding line */
            verticalSum.put(current.line,
                    (int)verticalSum.get(current.line) 
                    + (int)current.node.getValue());

            if(current.node.getLeft()!= null){
                //decrrement the line number when moving left.
                queue.add(new VerticalInfo(current.node.getLeft(), 
                                           current.line-1)
                );
            }
            if(current.node.getRight()!= null){
                // Increment the line when moving right
                queue.add(new VerticalInfo(current.node.getRight(), 
                                           current.line + 1)
                );
            }
        }
        return verticalSum;
    }
}

The complexity of the above implementation will be O(n) because we traverse all the nodes if the tree at least once.

Please share if there is anything wrong or missing. If you are looking for coaching to prepare for your interview, please reach book a free session