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.
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:
Time Complexity: O(N) where N is the number of nodes in given binary tree.
Auxiliary Space: O(h) where h is the height of binary tree due to recursion call stack.