Sum of all nodes with smaller values at a distance K from a given node in a BST
Given a Binary Search Tree, a target node in the BST, and an integer value K, the task is to find the sum of all nodes that are at a distance K from the target node whose value is less than the target node.
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.
Input: target = 7, K = 2
The nodes at a distance K(= 2) from the node 7 is 1, 4, and 6. Therefore, the sum of nodes is 11.
Input: target = 5, K = 1
Approach: The given problem can be solved by performing DFS Traversal for K distance below the target node and perform the DFS Traversal upward K distance from the target node. Follow the steps below to solve the problem:
- Define a function kDistanceDownSum(root, k, &sum) and perform the following steps:
- For the Base Case, check if the root is nullptr and k is less than 0, then return from the function.
- If the value of k equals 0, then add root->val to the variable sum and return.
- Call the same function kDistanceDownSum(root->left, k-1, sum) and kDistanceDownSum(root->right, k – 1, sum) for the left and right sub-trees.
- For the Base Case, check if the root is nullptr, then return -1.
- If the root is the same as the target, then call the function kDistanceDownSum(root->left, k – 1, sum) to calculate the sum for the first type of nodes and return 0(No second type of nodes possible).
- Initialize the variable dl as -1 and if the target is less than root, then set the value of dl as the value returned by the function kDistanceSum(root->left, target k, sum).
- If the value of dl is not equal to -1, then if sum equals (dl + 1), then add the value of root->data to the sum and then return -1.
- Similarly, initialize the variable dr as -1 and if the target is greater than the root, then update the value of dr to the value returned by kDistanceSum(root->right, target k, sum).
- If the value of dr is not equal to -1, then if the value of sum equals (dr + 1), then add the value of root->data to the sum. Otherwise, call the function kDistanceDownSum(root->left, k – dr – 2, sum) and return (1 + dr).
- After performing the above steps, print the value of ans as the resultant sum.
Following is the implementation of the above approach:
Auxiliary Space: O(1)