Sum of k smallest elements in BST

Given Binary Search Tree. The task is to find sum of all elements smaller than and equal to Kth smallest element.

Examples:

Input :  K = 3
              8
            /   \
           7     10
         /      /   \
        2      9     13
Output : 17
Explanation : Kth smallest element is 8 so sum of all
              element smaller then or equal to 8 are
              2 + 7 + 8

Input : K = 5
           8
         /   \
        5    11
      /  \
     2    7
      \
       3
Output :  25

Method 1 (Does not changes BST node structure)

The idea is to traverse BST in inorder traversal. Note that Inorder traversal of BST accesses elements in sorted (or increasing) order. While traversing, we keep track of count of visited Nodes and keep adding Nodes until the count becomes k.

// c++ program to find Sum Of All Elements smaller
// than or equal to Kth Smallest Element In BST
#include <bits/stdc++.h>
using namespace std;

/* Binary tree Node */
struct Node
{
    int data;
    Node* left, * right;
};

// utility function new Node of BST
struct Node *createNode(int data)
{
    Node * new_Node = new Node;
    new_Node->left = NULL;
    new_Node->right = NULL;
    new_Node->data = data;
    return new_Node;
}

// A utility function to insert a new Node
//  with given key in BST and also maintain lcount ,Sum
struct Node * insert(Node *root, int key)
{
    // If the tree is empty, return a new Node
    if (root == NULL)
        return createNode(key);

    // Otherwise, recur down the tree
    if (root->data > key)
        root->left = insert(root->left, key);

    else if (root->data < key)
        root->right = insert(root->right, key);

    // return the (unchanged) Node pointer
    return root;
}

// function return sum of all element smaller than
// and equal to Kth smallest element
int ksmallestElementSumRec(Node *root, int k, int &count)
{
    // Base cases
    if (root == NULL)
        return 0;
    if (count > k)
        return 0;

    // Compute sum of elements in left subtree
    int res = ksmallestElementSumRec(root->left, k, count);
    if (count >= k)
        return res;

    // Add root's data
    res += root->data;

    // Add current Node
    count++;
    if (count >= k)
      return res;

    // If count is less than k, return right subtree Nodes
    return res + ksmallestElementSumRec(root->right, k, count);
}

// Wrapper over ksmallestElementSumRec()
int ksmallestElementSum(struct Node *root, int k)
{
   int count = 0;
   ksmallestElementSumRec(root, k, count);
}

/* Driver program to test above functions */
int main()
{

    /*    20
        /    \
       8     22
     /   \
    4     12
         /   \
        10    14
          */
    Node *root = NULL;
    root = insert(root, 20);
    root = insert(root, 8);
    root = insert(root, 4);
    root = insert(root, 12);
    root = insert(root, 10);
    root = insert(root, 14);
    root = insert(root, 22);

    int k = 3;
    cout <<  ksmallestElementSum(root, k) <<endl;
    return 0;
}

Output :

22

Time complexity : O(k)

 

Method 2 (Efficient and changes structure of BST)

We can find the required sum in O(h) time where h is height of BST. Idea is similar to Kth-th smallest element in BST . Here we use augmented tree data structure to solve this problem efficiently in O(h) time [ h is height of BST ] .

Algorithm :

BST Node contain to extra fields : Lcount , Sum

For each Node of BST
lCount : store how many left child it has
Sum     : store sum of all left child it has

Find Kth smallest element
[ temp_sum store sum of all element less than equal to K ]

ksmallestElementSumRec(root, K, temp_sum)

  IF root -> lCount == K + 1
      temp_sum += root->data + root->sum;
      break;
  ELSE
     IF k > root->lCount   // Goto right sub-tree
        temp_sum += root->data + root-> sum;
        ksmallestElementSumRec(root->right, K-root->lcount+1, temp_sum)
     ELSE
        // Goto left sun-tree
        ksmallestElementSumRec( root->left, K, temp_sum)

Below is C++ implementation of above algo :

// C++ program to find Sum Of All Elements smaller
// than or equal t Kth Smallest Element In BST
#include <bits/stdc++.h>
using namespace std;

/* Binary tree Node */
struct Node
{
    int data;
    int lCount;
    int Sum ;
    Node* left;
    Node* right;
};

//utility function new Node of BST
struct Node *createNode(int data)
{
    Node * new_Node = new Node;
    new_Node->left = NULL;
    new_Node->right = NULL;
    new_Node->data = data;
    new_Node->lCount = 0 ;
    new_Node->Sum = 0;
    return new_Node;
}

// A utility function to insert a new Node with
// given key in BST and also maintain lcount ,Sum
struct Node * insert(Node *root, int key)
{
    // If the tree is empty, return a new Node
    if (root == NULL)
        return createNode(key);

    // Otherwise, recur down the tree
    if (root->data > key)
    {
        // increment lCount of current Node
        root->lCount++;

        // increment current Node sum by adding
        // key into it
        root->Sum += key;

        root->left= insert(root->left , key);
    }
    else if (root->data < key )
        root->right= insert (root->right , key );

    // return the (unchanged) Node pointer
    return root;
}

// function return sum of all element smaller than and equal
// to Kth smallest element
void ksmallestElementSumRec(Node *root, int k , int &temp_sum)
{
    if (root == NULL)
        return ;

    // if we fount k smallest element then break the function
    if ((root->lCount + 1) == k)
    {
        temp_sum += root->data + root->Sum ;
        return ;
    }

    else if (k > root->lCount)
    {
        // store sum of all element smaller than current root ;
        temp_sum += root->data + root->Sum;

        // decremented k and call right sub-tree
        k = k -( root->lCount + 1);
        ksmallestElementSumRec(root->right , k , temp_sum);
    }
    else // call left sub-tree
        ksmallestElementSumRec(root->left , k , temp_sum );
}

// Wrapper over ksmallestElementSumRec()
int ksmallestElementSum(struct Node *root, int k)
{
    int sum = 0;
    ksmallestElementSumRec(root, k, sum);
    return sum;
}

/* Driver program to test above functions */
int main()
{
    /*    20
        /    \
       8     22
     /   \
    4     12
         /   \
        10    14
          */
    Node *root = NULL;
    root = insert(root, 20);
    root = insert(root, 8);
    root = insert(root, 4);
    root = insert(root, 12);
    root = insert(root, 10);
    root = insert(root, 14);
    root = insert(root, 22);

    int k = 3;
    cout <<  ksmallestElementSum(root, k) << endl;
    return 0;
}

Output:

 22

Time Complexity : O(h) where h is height of tree.
This article is contributed by Nishant Singh. If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please write comments if you find anything incorrect, or you want to share more information about the topic discussed above.

GATE CS Corner    Company Wise Coding Practice

Recommended Posts:







Writing code in comment? Please use ide.geeksforgeeks.org, generate link and share the link here.