I am going to write the full implementation for the problem in order to make it easy to prove my arguments about the time being taken.
.
Assuming that each node of BST has the following structure:
typedef struct node{
int vale;
struct node* left;
struct node* right;
}node;
The algorithm will have 2 steps:
a. Start from the root node of the Tree and find the starting node and all the ancestor of that node. Store all this in the stack passed:
//root -> root node of the tree.
//val -> value of the node to find.
// s -> stack to store all ancestor.
node* find_node(node* root, int val,std::stack<node*> &s)
{
if(root != NULL) s.push(root);
if(root == NULL || root->value == val) return root;
if(root->value > val) return find_node(root->left);
else return find_node(root->right);
}
and the call to this method would look like:
//Assuming that the root is the root node of the tree.
std::stack<node*> s;
node* ptr = find_node(root,s); // we have all ancestors of ptr along with ptr in stack s now.
b. Now we have to print next k immediate bigger (than ptr) elements of the tree. We will start from the node (which is ptr) :
// s -> stack of all ancestor of the node.
// k -> k-successor to find.
void print_k_bigger(stack<node*> s, int k)
{
//since the top element of stack is the starting node. So we won't print it.
// We will just pop the first node and perform inorder traversal on its right child.
prev = s.top();
s.pop();
inorder(prev->right,k);
// Now all the nodes present in the stack are ancestor of the node.
while(!s.empty() && k>0)
{
//pop the node at the top of stack.
ptr = s.top();
s.pop();
//if the node popped previously (prev) was the right child of the current
//node, i.e. prev was bigger than the current node. So we will have to go
//one more level up to search for successors.
if(prev == ptr->right) continue;
//else the current node is immidiate bigger than prev node. Print it.
printf("%d",ptr->value);
//reduce the count.
k--;
//inorder the right subtree of the current node.
inorder(ptr->right);
//Note this.
prev = ptr;
}
}
Here is how our inorder will look like:
void inorder(node* ptr, int& k)
{
if(ptr != NULL)
{
inorder(ptr->left,k);
printf("%d",ptr->value);
k--;
inorder(ptr->right,k);
}
}
Time Analysis:
- The method
find_node is O(h) as it can go upto the length max root to leaf path.
- The method
print_k_bigger is O(h+k) as in every iteration of the loop, either the size of stack is reducing or the value of k is reducing. Note that when inorder() is called from inside the while loop, it won't raise additional overhead as all the calls to inorder() together will take at max O(k).
