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:
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]
- Nodes at width position 0: 0, 2. Therefore, the sum of nodes is 1 + 3 = 4.
- 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]
- Nodes at width position -1 is 3. Therefore, the sum of nodes is 4.
- Nodes at width position 0: 0, 2. Therefore, the sum of nodes is 1 + 3 = 4.
- 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 nodestruct Node { int data; Node* left; Node* right;};// Function to create new nodeNode* 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 positionvoid 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 queriesvoid 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 Codeint 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 approachimport java.io.*;import java.util.*;class GFG{ // Structure of a nodestatic class Node{ int data; Node left; Node right;}// Function to create new nodestatic 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 positionstatic 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 queriesstatic 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 Codepublic 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. |
Python3
# C++ program for the above approachfrom collections import defaultdict# Structure of a nodeclass Node: def __init__(self, data): self.data = data self.left = None self.right = None# Function to create new nodedef newNode(d): n = Node(d) n.left = None n.right = None return n# Function to pre-compute the sum of# weights at each width positiondef findwt(root, wt, um, width=0): if root is None: return um[width] += wt[root.data] findwt(root.left, wt, um, width-1) findwt(root.right, wt, um, width+1)def solveQueries(wt, root, queries): um = defaultdict(int) findwt(root, wt, um) x = 0 for k, v in sorted(um.items()): x += v um[k] = x for q in queries: l = q[0] r = q[1] if l - 1 in um: l_val = um[l - 1] else: l_val = 0 if r in um: r_val = um[r] else: r_val = 0 print(r_val - l_val)if __name__ == '__main__': N = 8 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) wt = [8, 6, 4, 5, 1, 2, 9, 1] queries = [[-1, 1], [-2, -1], [0, 3]] solveQueries(wt, root, queries) |
C#
using System;using System.Collections.Generic;class GFG { // Structure of a node public class Node { public int data; public Node left; public 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, SortedDictionary<int, int> um, int width) { // Base Case if (root == null) { return; } // Update the current width // position weight if (um.ContainsKey(width)) { um[width] = um[width] + wt[root.data]; } else { um.Add(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 SortedDictionary<int, int> um = new SortedDictionary<int, int>(); // Function Call to fill um findwt(root, wt, um, 0); // Stores the sum of all previous // nodes, while traversing Map int x = 0; // Create new TreeMap to store modified values SortedDictionary<int, int> newUm = new SortedDictionary<int, int>(); // Traverse the Map um foreach(KeyValuePair<int, int> it in um) { x += it.Value; newUm.Add(it.Key, x); } // Replace the original TreeMap with the new one um = newUm; // 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[r]; } if (um.ContainsKey(l - 1)) { ans -= um[l - 1]; } Console.WriteLine(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 = { new int[] { -1, 1 }, new int[] { -2, -1 }, new int[] { 0, 3 } }; solveQueries(wt, root, queries); }} |
Javascript
// Structure of a nodeclass Node { constructor(data) { this.data = data; this.left = null; this.right = null; }}// Function to create new nodefunction newNode(d) { const n = new Node(d); n.left = null; n.right = null; return n;}// Function to pre-compute the sum of// weights at each width positionfunction findwt(root, wt, um, width=0) { // Base Case if (root == null) { return; } // Update the current width position weight if (um.has(width)) { um.set(width, um.get(width) + wt[root.data]); } else { um.set(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 queriesfunction solveQueries(wt, root, queries) { // Stores the weight sum of each width position const um = new Map(); // Function Call to fill um findwt(root, wt, um, 0); // Stores the sum of all previous nodes, while traversing Map let x = 0; // Traverse the Map um for (const [key, value] of um) { x += value; um.set(key, x); } // Iterate over all queries for (let i = 0; i < queries.length; i++) { const l = queries[i][0]; const r = queries[i][1]; // Print the result for the current query [l, r] let ans = 0; if (um.has(r)) { ans = um.get(r); } if (um.has(l - 1)) { ans -= um.get(l - 1); } console.log(ans); }}// Driver Code(function main() { const N = 8; // Given Tree const 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); const wt = [8, 6, 4, 5, 1, 2, 9, 1]; const queries = [[-1, 1], [-2, -1], [0, 3]]; solveQueries(wt, root, queries);})(); |
Output
25 7 29
Time Complexity: O(N*log N + Q)
Auxiliary Space: O(N)
Feeling lost in the world of random DSA topics, wasting time without progress? It’s time for a change! Join our DSA course, where we’ll guide you on an exciting journey to master DSA efficiently and on schedule.
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
