Median of All Nodes from a Given Range in a BST

Firdausia Fatima
Last Updated: May 13, 2022

Introduction

Tree traversal, also known as walking the tree, refers to the process of visiting each node in a tree data structure exactly once. These traversals are classified based on the order in which the nodes are traversed. Preorder, in order, and postorder are the three types of traversal. We'll talk about a problem that can be solved with a modification to traverse today.

Problem Statement

Given a binary search tree ‘ROOT’ with ‘N’ nodes and two integers ‘L’ and ‘R’. You have to find the median of all the nodes in the given BST whose value range from ‘L’ to ‘R’.

Example

Approach

In a Binary Search Tree, if we perform an inorder traversal and save the node's value in a vector, it will be sorted. To get all the nodes whose values range from 'L' to 'R,' we'll use a modified inorder traversal. What we’ll do is maintain three states, first one indicates that we’ve not entered the range(VAL < L), the second indicates that we are in the range(VAL >= L && VAL <= R), and the third one indicates that we’ve passed the range (VAL > R).

Therefore, if the size of the vector is odd, we'll return the middle element; otherwise, we'll return the average of the middle two elements.

We'll be using two functions: findMedian(), which finds the median of all nodes from a given range in a BST, and getNodesInRange(), which returns an array of nodes in a given range.

Let’s see how these functions work.

findMedian()

Parameters

  1. ‘ROOT’ - Pointer to Binary Search Tree Node.
  2. ‘L’ - Start of Range.
  3. ‘R’ - End of the Range.

Working

  1. Check if the tree is empty, i.e., ROOT = NULL, then return -1.
  2. Next, declare a vector, ‘RANGE_NODES’ and call the getNodesInRange() function to get all the node’s values in range ‘L’ to ‘R’.
  3. If the size of the 'RANGE_NODES' vector is odd, then return the middle element, else return the average of the middle two elements.

getNodesInRange()

Parameters:

  1. ‘ROOT’ - Pointer to Binary Search Tree Node.
  2. ‘L’ - Start of Range.
  3. ‘R’ - End of the Range.
  4. ‘RANGE_NODES’ - Array to store the nodes whose value is in the range [‘L’, ‘R’].
  5. 'FLAG' will have three values 0, 1, and 2.
    1. If FLAG = 0, we've not entered the range.
    2. If FLAG = 1, then we are in the range.
    3. If FLAG = 2, then we are done traversing.

Working:

  1. Base Case - If ROOT = NULL or FLAG = 2, we know we are done traversing hence return.
  2. Recurse in the left subtree.
  3. Check if the current node's value is greater than equal to 'L' the update FLAG = 1.
  4. Check if we exceed the range, update FLAG = 2.
  5. If we are in the range, push back the current node value into 'RANGE_NODES.'
  6. Recurse in the right subtree.

Program

#include <iostream>
#include <vector>
using namespace std;

// Class for Tree Node.
class TreeNode
{
public:
  int data;
  // Pointer to the left node.
  TreeNode *left;
  // Pointer to the right node.
  TreeNode *right;
  TreeNode(int data)
  {
      this->data = data;
      this->left = this->right = NULL;
  }
};

// Function to get all nodes of BST whose value lie in range [L, R].
void getNodesInRange(TreeNode *root, int l, int r, vector<int> &rangeNodes, int &flag)
{
  // If 'ROOT' is NULL or we've reached the end of the range, return null.
  if (root == NULL || flag == 2)
      return;
      
  // Traverse the left subtree.
  getNodesInRange(root->left, l, r, rangeNodes, flag);
  
  // Check if we've not entered the range, i.e., FLAG = 0, and the current node value is greater than L, then update FLAG = 1.
  if (flag == 0 && root->data >= l)
  {
      flag = 1;
  }
  
  // Check if we've exceeded the range, update FLAG=2.
  if (flag == 1 && root->data > r)
  {
      flag = 2;
  }
  
  // If we are an inside range [L, R], then push this value into the 'RANGE_NODES' array.
  if (flag == 1)
  {
      rangeNodes.emplace_back(root->data);
  }
  
  // Traverse right subtree.
  getNodesInRange(root->right, l, r, rangeNodes, flag);
}

// Function that returns the Median of all nodes from a given range in a BST.
double findMedian(TreeNode *root, int l, int r)
{
  // If BST is empty, return.
  if (root == NULL)
      return -1;
      
  // Declare a vector to store the nodes in range [L, R].
  vector<int> rangeNodes;
  int flag = 0;
  
  // Call 'getNodesInRange()' Function to get all the nodes in BST whose value lies in the range [L, R].
  getNodesInRange(root, l, r, rangeNodes, flag);
  
  // Get the size.
  int size = rangeNodes.size();
  
  // If there are no elements in range return -1.
  if (size == 0)
      return -1;
      
  // If the number of nodes is Odd, return the middle node value.
  if (size & 1)
  {
      return (double)rangeNodes[size / 2];
  }
  
  // Else return the average of two middle nodes.
  return ((double)rangeNodes[size / 2 - 1] + (double)rangeNodes[size / 2]) / 2;
}

TreeNode *insert(TreeNode *root, int data)
{
  if (!root)
  {
      return new TreeNode(data);
  }
  if (data > root->data)
  {
      root->right = insert(root->right, data);
  }
  else
  {
      root->left = insert(root->left, data);
  }
  return root;
}

// Function to create a BST.
TreeNode *createBST(vector<int> &nodes)
{
  TreeNode *root = NULL;
  for (int &node : nodes)
  {
      root = insert(root, node);
  }
  return root;
}

// Main.
int main()
{
  int n;
  cout << "Enter the number of elements in BST: ";
  cin >> n;
  
  vector<int> nodes(n);
  cout << "Enter the value of nodes in level order fashion: ";
  for (int i = 0; i < n; i++)
  {
      cin >> nodes[i];
  }
  
  TreeNode *root = createBST(nodes);
  cout << "Enter the values of L and R: ";
  int l, r;
  cin >> l >> r;
  
  double median = findMedian(root, l, r);
  
  cout << "The median of all nodes in the range [" << l << ", " << r << "] is: " << median;
  return 0;
}

Input

Enter the number of elements in BST: 7
Enter the values of nodes in level order fashion: 7 4 9 3 8 13 11
Enter the values of L and R: 5 12

Output

The median of all nodes in the range [5, 12] is: 8.5

Time Complexity

O(N), where ‘N’ is the number of nodes in the binary search tree.

In the function getNodesInRange(), we transverse all ‘N’ nodes in the worst case. Hence, the time complexity is O(N).

Space Complexity

O(N), where ‘N’ is the number of nodes in the binary search tree.

Since we are declaring the 'RANGE_NODES' vector, which can grow up to a size of 'N', space complexity is O(N).

Key Takeaways

Tree Traversal is the most basic algorithm to know when it comes to problems on Trees. Other Tree algorithms include BFSDFS, and Binary Heaps.

Head over to CodeStudio to read other such exciting blogs on Trees. 

Recently Coding Ninjas has released a specially designed test series for acing Interviews- CodeStudio Test Series.

Thanks for reading. I hope you’ve gained a lot from this blog.

Was this article helpful ?
0 upvotes