Red-black tree access by ordinal index

Question

I have a red-black tree (binary tree, all the leaves are within 2 levels). I can navigate through nodes: to left, right or parent. I know the entire count of nodes.

I have to find the N-th smallest element in a tree. Is there any way to do this faster than in O(n)? Any ideas of optimizing access by index?

Answer 1

In each node X you should store, how many nodes are in subtree with X as a root.

count(LEAF) = 1
count(NODE) = count(NODE->LEFT) + count(NODE->RIGHT) + 1

During each insert/delete you should use this equation to update counts in nodes affected by rotations.

After that solution is simple

NODE nth(NODE root, int n) {
    if (root->left->count <= n) return nth(root->left, n);
    if ( root->left->count + 1 == n) return root;
    return nth(root->right, n - root->left->count - 1);
}