Diagonal Sum of a Binary Tree

Consider lines of slope -1 passing between nodes (dotted lines in below diagram). Diagonal sum in a binary tree is sum of all node’s data lying between these lines. Given a Binary Tree, print all diagonal sums.

For the following input tree, output should be 9, 19, 42.
9 is sum of 1, 3 and 5.
19 is sum of 2, 6, 4 and 7.
42 is sum of 9, 10, 11 and 12.

Recommended: Please try your approach on {IDE} first, before moving on to the solution.

Algorithm:
The idea is to keep track of vertical distance from top diagonal passing through root. We increment the vertical distance we go down to next diagonal.
1. Add root with vertical distance as 0 to the queue.
2. Process the sum of all right child and right of right child and so on.
3. Add left child current node into the queue for later processing. The vertical distance of left child is vertical distance of current node plus 1.
4. Keep doing 2nd, 3rd and 4th step till the queue is empty.

Following is the implementation of above idea.

C++

// C++ Program to find diagonal 
// sum in a Binary Tree 
#include <bits/stdc++.h> 
using namespace std; 
  
struct Node 
{ 
    int data; 
    struct Node* left; 
    struct Node* right; 
}; 
  
struct Node* newNode(int data) 
{ 
    struct Node* Node = 
            (struct Node*)malloc(sizeof(struct Node)); 
      
    Node->data = data; 
    Node->left = Node->right = NULL; 
  
    return Node; 
} 
  
// root - root of the binary tree 
// vd - vertical distance diagonally 
// diagonalSum - map to store Diagonal  
// Sum(Passed by Reference) 
void diagonalSumUtil(struct Node* root, 
                int vd, map<int, int> &diagonalSum) 
{ 
    if(!root) 
        return; 
          
    diagonalSum[vd] += root->data; 
  
    // increase the vertical distance if left child 
    diagonalSumUtil(root->left, vd + 1, diagonalSum); 
  
    // vertical distance remains same for right child 
    diagonalSumUtil(root->right, vd, diagonalSum); 
} 
  
// Function to calculate diagonal  
// sum of given binary tree 
void diagonalSum(struct Node* root) 
{ 
  
    // create a map to store Diagonal Sum 
    map<int, int> diagonalSum;  
      
    diagonalSumUtil(root, 0, diagonalSum); 
  
    map<int, int>::iterator it; 
        cout << "Diagonal sum in a binary tree is - "; 
      
    for(it = diagonalSum.begin(); 
                it != diagonalSum.end(); ++it) 
    { 
        cout << it->second << " "; 
    } 
} 
  
// Driver code 
int main() 
{ 
    struct Node* root = newNode(1); 
    root->left = newNode(2); 
    root->right = newNode(3); 
    root->left->left = newNode(9); 
    root->left->right = newNode(6); 
    root->right->left = newNode(4); 
    root->right->right = newNode(5); 
    root->right->left->right = newNode(7); 
    root->right->left->left = newNode(12); 
    root->left->right->left = newNode(11); 
    root->left->left->right = newNode(10); 
  
    diagonalSum(root); 
  
    return 0; 
} 
  
// This code is contributed by Aditya Goel 

Java

// Java Program to find diagonal sum in a Binary Tree 
import java.util.*; 
import java.util.Map.Entry; 
  
//Tree node 
class TreeNode 
{ 
    int data; //node data 
    int vd; //vertical distance diagonally 
    TreeNode left, right; //left and right child's reference 
  
    // Tree node constructor 
    public TreeNode(int data) 
    { 
        this.data = data; 
        vd = Integer.MAX_VALUE; 
        left = right = null; 
    } 
} 
  
// Tree class 
class Tree 
{ 
    TreeNode root;//Tree root 
  
    // Tree constructor 
    public Tree(TreeNode root)  {  this.root = root;  } 
  
    // Diagonal sum method 
    public void diagonalSum() 
    { 
        // Queue which stores tree nodes 
        Queue<TreeNode> queue = new LinkedList<TreeNode>(); 
  
        // Map to store sum of node's data lying diagonally 
        Map<Integer, Integer> map = new TreeMap<>(); 
  
        // Assign the root's vertical distance as 0. 
        root.vd = 0; 
  
        // Add root node to the queue 
        queue.add(root); 
  
        // Loop while the queue is not empty 
        while (!queue.isEmpty()) 
        { 
            // Remove the front tree node from queue. 
            TreeNode curr = queue.remove(); 
  
            // Get the vertical distance of the dequeued node. 
            int vd = curr.vd; 
  
            // Sum over this node's right-child, right-of-right-child 
            // and so on 
            while (curr != null) 
            { 
                int prevSum = (map.get(vd) == null)? 0: map.get(vd); 
                map.put(vd, prevSum + curr.data); 
  
                // If for any node the left child is not null add 
                // it to the queue for future processing. 
                if (curr.left != null) 
                { 
                    curr.left.vd = vd+1; 
                    queue.add(curr.left); 
                } 
  
                // Move to the current node's right child. 
                curr = curr.right; 
            } 
        } 
  
        // Make an entry set from map. 
        Set<Entry<Integer, Integer>> set = map.entrySet(); 
  
        // Make an iterator 
        Iterator<Entry<Integer, Integer>> iterator = set.iterator(); 
  
        // Traverse the map elements using the iterator. 
         System.out.print("Diagonal sum in a binary tree is - "); 
        while (iterator.hasNext()) 
        { 
            Map.Entry<Integer, Integer> me = iterator.next(); 
  
            System.out.print(me.getValue()+" "); 
        } 
    } 
} 
  
//Driver class 
public class DiagonalSum 
{ 
    public static void main(String[] args) 
    { 
        TreeNode root = new TreeNode(1); 
        root.left = new TreeNode(2); 
        root.right = new TreeNode(3); 
        root.left.left = new TreeNode(9); 
        root.left.right = new TreeNode(6); 
        root.right.left = new TreeNode(4); 
        root.right.right = new TreeNode(5); 
        root.right.left.left = new TreeNode(12); 
        root.right.left.right = new TreeNode(7); 
        root.left.right.left = new TreeNode(11); 
        root.left.left.right = new TreeNode(10); 
        Tree tree = new Tree(root); 
        tree.diagonalSum(); 
    } 
} 

Python3

# Program to find diagonal sum in a Binary Tree 
  
class newNode:  
    def __init__(self, data):  
        self.data = data  
        self.left = self.right = None
          
# Function to compute height and  
# root - root of the binary tree  
# vd - vertical distance diagonally  
# diagonalSum - map to store Diagonal  
# Sum(Passed by Reference)  
def diagonalSumUtil(root, vd, diagonalSum) : 
  
    if(not root):  
        return
          
    if vd not in diagonalSum: 
        diagonalSum[vd] = 0
    diagonalSum[vd] += root.data  
  
    # increase the vertical distance 
    # if left child  
    diagonalSumUtil(root.left, vd + 1,  
                          diagonalSum)  
  
    # vertical distance remains same  
    # for right child  
    diagonalSumUtil(root.right, vd, 
                       diagonalSum)  
  
# Function to calculate diagonal  
# sum of given binary tree  
def diagonalSum(root) : 
  
    # create a map to store Diagonal Sum  
    diagonalSum = dict()  
      
    diagonalSumUtil(root, 0, diagonalSum)  
  
    print("Diagonal sum in a binary tree is - ",  
                                       end = "") 
      
    for it in diagonalSum: 
        print(diagonalSum[it], end = " ") 
          
# Driver Code  
if __name__ == '__main__': 
    root = newNode(1)  
    root.left = newNode(2)  
    root.right = newNode(3)  
    root.left.left = newNode(9)  
    root.left.right = newNode(6)  
    root.right.left = newNode(4)  
    root.right.right = newNode(5)  
    root.right.left.right = newNode(7)  
    root.right.left.left = newNode(12)  
    root.left.right.left = newNode(11)  
    root.left.left.right = newNode(10)  
  
    diagonalSum(root) 
  
# This code is contributed  
# by SHUBHAMSINGH10 

Output:

Diagonal sum in a binary tree is - 9 19 42

Exercise:
This problem was for diagonals from top to bottom and slope -1. Try the same problem for slope +1.

This article is contributed by Kumar Gautam. Please write comments if you find anything incorrect, or you want to share more information about the topic discussed above.

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.

My Personal Notes

Recommended: Please try your approach on {IDE} first, before moving on to the solution.

C++

Java

Python3

Recommended Posts: