Vertical Sum in a given Binary Tree

Question

Given a Binary Tree, find vertical sum of the nodes that are in same vertical line. Print all sums through different vertical lines.

To understand what's same vertical line, we need to define horizontal distances first. If two nodes have the same Horizontal Distance (HD), then they are on same vertical line. The idea of HD is simple. HD for root is 0, a right edge (edge connecting to right subtree) is considered as +1 horizontal distance and a left edge is considered as -1 horizontal distance. For example, in the above tree, HD for Node 4 is at -2, HD for Node 2 is -1, HD for 5 and 6 is 0 and HD for node 7 is +2.

Examples:

    1
  /   \
 2     3
/ \   / \
4  5  6  7

The tree has 5 vertical lines

Vertical-Line-1 has only one node 4 => vertical sum is 4

Vertical-Line-2: has only one node 2=> vertical sum is 2

Vertical-Line-3: has three nodes: 1,5,6 => vertical sum is 1+5+6 = 12

Vertical-Line-4: has only one node 3 => vertical sum is 3

Vertical-Line-5: has only one node 7 => vertical sum is 7

So expected output is 4, 2, 12, 3 and 7

My solution: I think out a o(nlong(n)) solution for this problem. The idea is:

(1) using preorder traversal to get the HD for every node, and store the HD and its associated node in an array.

(2) sort the array by HD

(3) traversal the sorted array to print result.

I'm sure this is not the best one for this problem. Can anyone help me give a better solution?

Answer 1

Can't you do it all in the first traversal? Define a dictionary (hash map) from HD to sum. And for each node you visit add its value to the right dictionary key - this is a O(n) solution.

d = {}

def traverse(node, hd):
  if not node:
    return
  if not hd in d:
    d[hd] = 0
  d[hd] = d[hd] + node.value
  traverse(node.left, hd - 1)
  traverse(node.right, hd + 1)

Then just call traverse(root, 0)