Skip to content
Related Articles

Related Articles

Improve Article

Queries to find the sum of weights of all nodes with vertical width from given range in a Binary Tree

  • Last Updated : 01 Jun, 2021

Given a Binary Tree consisting of N nodes having values in the range [0, N – 1] rooted as 0, an array wt[] of size N where wt[i] is the weight of the ith node and a 2D array Q[][] consisting of queries of the type {L, R}, the task for each query is to find the sum of weights of all nodes whose vertical width lies in the range [L, R].

Examples:

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.  To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

In case you wish to attend live classes with experts, please refer DSA Live Classes for Working Professionals and Competitive Programming Live for Students.

Input: N = 4, wt[] = {1, 2, 3, 4}, Q[][] = {{0, 1}, {1, 1}}



                                 0(1)
                                /     \
                          3(4)    1(2)
                                    /  
                               2(3)
Output:
6
10
Explanation:
For Query 1: The range is [0, 1]

  1. Nodes at width position 0: 0, 2. Therefore, the sum of nodes is 1 + 3 = 4.
  2. Nodes at width position 1: 1. Therefore, the sum of nodes is 2.

Therefore, the sum of weights all nodes = 4 + 2 = 6.

For Query 2: The range is [-1, 1]

  1. Nodes at width position -1 is 3. Therefore, the sum of nodes is 4.
  2. Nodes at width position 0: 0, 2. Therefore, the sum of nodes is 1 + 3 = 4.
  3. Nodes at width position 1: 1. Therefore, the sum of nodes is 2.

Therefore, the sum of weights all nodes = 4 + 4 + 2 = 10.

Input: N= 8, wt[] = {8, 6, 4, 5, 1, 2, 9, 1}, Q[][] = {{-1, 1}, {-2, -1}, {0, 3}}

                                           1
                                         /  \
                                       3     2
                                    /   \      \
                                  5     6     4
                                               /  \ 
                                             7    0
Output:
25
7
29

Naive Approach: The simplest approach to solve the given problem is to perform a traversal on the tree for each query and then print the sum of weights of all the nodes that lie in the given range [L, R].

Time Complexity: O(N*Q)
Auxiliary Space: O(1)

Efficient Approach: The above approach can also be optimized by precomputing the sum of weights for every width position and store them in a map M, then find the prefix sum of the newly created array to process the queries in constant time. Follow the below steps to solve the problem:

  • Perform the DFS traversal on the given tree and update the sum at all the width positions and store them in a map M by performing the following operations:
    • Initially, the position pos of the root is 0.
    • In each recursive call, update the weight of the node in M by updating M[pos] = M[pos] + wt[node].
    • Recursively call to its left child by decrementing pos by 1.
    • Recursively call to its right child by incrementing pos by 1.
  • Iterate over the map M and update the value of each key-value pair to the sum of the previous values that occurred.
  • Now, traverse the array Queries[] and for each query [L, R], print the value of (M[R] – M[L – 1]) as the result.

Below is the implementation of the above approach:

C++




// C++ program for the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Structure of a node
struct Node {
    int data;
    Node* left;
    Node* right;
};
 
// Function to create new node
Node* newNode(int d)
{
    Node* n = new Node;
    n->data = d;
    n->left = NULL;
    n->right = NULL;
    return n;
}
 
// Function to pre-compute the sum of
// weights at each width position
void findwt(Node* root, int wt[],
            map<int, int>& um,
            int width = 0)
{
    // Base Case
    if (root == NULL) {
        return;
    }
 
    // Update the current width
    // position weight
    um[width] += wt[root->data];
 
    // Recursive Call to its left
    findwt(root->left, wt, um,
           width - 1);
 
    // Recursive Call to its right
    findwt(root->right, wt, um,
           width + 1);
}
 
// Function to find the sum of the
// weights of nodes whose vertical
// widths lies over the range [L, R]
// for Q queries
void solveQueries(int wt[], Node* root,
                  vector<vector<int> > queries)
{
    // Stores the weight sum
    // of each width position
    map<int, int> um;
 
    // Function Call to fill um
    findwt(root, wt, um);
 
    // Stores the sum of all previous
    // nodes, while traversing Map
    int x = 0;
 
    // Traverse the Map um
    for (auto it = um.begin();
         it != um.end(); it++) {
        x += it->second;
        um[it->first] = x;
    }
 
    // Iterate over all queries
    for (int i = 0;
         i < queries.size(); i++) {
 
        int l = queries[i][0];
        int r = queries[i][1];
 
        // Print the result for the
        // current query [l, r]
        cout << um[r] - um[l - 1]
             << "\n";
    }
}
 
// Driver Code
int main()
{
    int N = 8;
 
    // Given Tree
    Node* root = newNode(1);
    root->left = newNode(3);
    root->left->left = newNode(5);
    root->left->right = newNode(6);
    root->right = newNode(2);
    root->right->right = newNode(4);
    root->right->right->left = newNode(7);
    root->right->right->right = newNode(0);
    int wt[] = { 8, 6, 4, 5, 1, 2, 9, 1 };
    vector<vector<int> > queries{ { -1, 1 },
                                  { -2, -1 },
                                  { 0, 3 } };
 
    solveQueries(wt, root, queries);
 
    return 0;
}

Java




// Java program for the above approach
import java.io.*;
import java.util.*;
 
class GFG{
   
// Structure of a node
static class Node
{
    int data;
    Node left;
    Node right;
}
 
// Function to create new node
static Node newNode(int d)
{
    Node n = new Node();
    n.data = d;
    n.left = null;
    n.right = null;
    return n;
}
 
// Function to pre-compute the sum of
// weights at each width position
static void findwt(Node root, int wt[],
                   TreeMap<Integer, Integer> um,
                   int width)
{
     
    // Base Case
    if (root == null)
    {
        return;
    }
     
    // Update the current width
    // position weight
    if (um.containsKey(width))
    {
        um.put(width, um.get(width) +
               wt[root.data]);
    }
    else
    {
        um.put(width, wt[root.data]);
    }
 
    // Recursive Call to its left
    findwt(root.left, wt, um,
           width - 1);
 
    // Recursive Call to its right
    findwt(root.right, wt, um,
           width + 1);
}
 
// Function to find the sum of the
// weights of nodes whose vertical
// widths lies over the range [L, R]
// for Q queries
static void solveQueries(int wt[], Node root,
                         int[][] queries)
{
     
    // Stores the weight sum
    // of each width position
    TreeMap<Integer,
            Integer> um = new TreeMap<Integer,
                                      Integer>();
 
    // Function Call to fill um
    findwt(root, wt, um, 0);
 
    // Stores the sum of all previous
    // nodes, while traversing Map
    int x = 0;
 
    // Traverse the Map um
      for(Map.Entry<Integer, Integer> it : um.entrySet())
      {
        x += it.getValue();
          um.put(it.getKey(), x);
    }
 
    // Iterate over all queries
    for(int i = 0; i < queries.length; i++)
    {
        int l = queries[i][0];
        int r = queries[i][1];
 
        // Print the result for the
        // current query [l, r]
          int ans = 0;
          if (um.containsKey(r))
          {
              ans = um.get(r);
        }
           
          if (um.containsKey(l - 1))
          {
              ans -= um.get(l - 1);
        }
        System.out.println(ans);
    }
}
 
// Driver Code
public static void main(String[] args)
{
    int N = 8;
 
    // Given Tree
    Node root = newNode(1);
    root.left = newNode(3);
    root.left.left = newNode(5);
    root.left.right = newNode(6);
    root.right = newNode(2);
    root.right.right = newNode(4);
    root.right.right.left = newNode(7);
    root.right.right.right = newNode(0);
     
    int wt[] = { 8, 6, 4, 5, 1, 2, 9, 1 };
    int queries[][] = { { -1, 1 },
                        { -2, -1 },
                        { 0, 3 } };
 
    solveQueries(wt, root, queries);
}
}
 
// This code is contributed by Dharanendra L V.
Output: 
25
7
29

 

Time Complexity: O(N*log N + Q)
Auxiliary Space: O(N)




My Personal Notes arrow_drop_up
Recommended Articles
Page :