Algorithm for determining the number of nodes with a key greater than an integer K on a BST

Question

I had the following problem in a test a week ago. I haven't gotten my grade back, but I'm sure my solution didn't fully target all the base cases of the problem.

The statement is the following:

For a Binary Searcht tree, Write an algorithm (using Pseudo-code) that computes the number of nodes with a key greater or equal to a given integer k. Your algorithm should run in the worst-case time O(h), where h is the height of the Binary Search Tree.

Assume you're given a method subtreeSize(treeNode n) that runs in time O(1), and returns the number of nodes in the subtree rooted at n, including n itself.

This is my solution:

nbNodesGreaterEqual(treeNode n, int k){

    if(n == null) return 0;
    if(n.getValue() >= k) return 1 + substreeSize(n.getRightChild()) + nbNodesGreaterEqual(n.getLeftChild(), k);
    if(n.getValue < k) return nbNodesGreaterEqual(n.getRightChild,k);
}

Is my algorithm complete? Also, is there a way to write this same algorithm for a regular binary tree (not a BST) that doesn't traverse through all the nodes?

Answer 1

Assuming right child of a node holds a key larger or equal than the key stored in that node. The idea is to go down the right branch until we find a node whose key is larger or equal to k. If we find a node with key equal to k, we know that all nodes on the right should have keys larger to k. So we take the size of that subtree. Add 1 with that because this nodes also needs to be counted. In case we stopped going down the right branch when found a node with key larger than k, we know that the right subtree needs to be included and there is possibility of some acceptable nodes being found in the left subtree. So we call the function recursively on the left subtree. Finally, 1 is added because this node should be counted too. Base case is when node is null, we return 0.

ksum(node, k)      
  x = 0
  node==null
      return 0
  if node->key < k
     x = ksum(node->right, k)
  else if node->key == k
     x = subtreesize(node->right) + 1
  else
     x = subtreesize(node->right) + ksum(node->left, k) + 1

  return x