Find K-th next element in O(log n) time

Question

We want to store a set S of distinct positive integers and support the following operations on S:

Insert(a): insert integer a !∈ S into S.

Delete(a): delete integer a ∈ S from S.

FindNext(a, k), where a ∈ S and k is an integer: if S consists of some elements a1 < a2 < · · · < an, and a = ai (for some i, 1 ≤ i ≤ n), then return ai+k (for simplicity, assume that i + k ≤ n).

The worst-case time complexity of each operation should be O(log n) where n is the size of S. An augmented data structure can be used (without any modification) “as a black box” to achieve the above goal: using the operations that this data structure provides, the algorithms for performing the above three operations are quite simple. Give an algorithm for the FindNext(a, k) operation that uses only the operations of this data structure and consists of at most 5 lines of pseudo-code.

You are allowed to make use of other operations available with this data structure, i.e. everything you can make run in the specified time.

Answer 1

You can use the order-statistic tree which is a red-black tree augmented each node $x$ with an additional field $size(x)$. The field contains the number of nodes in the subtree rooted at $x$ (including $x$ itself).

As red-black tree, order-statistic tree supports Query, Insert, and Delete in $O(\lg n)$ time. Further, order-statistic tree provides two new operations:

1. Select(i): find the element with rank $i$ in $O(\lg n)$ time.
2. Rank(x): obtain the rank of element $x$ in $O(\lg n)$ time.

Then FindNext(a, k) can be implemented as Select( Rank(x) + k ). The time complexity is obviously $O(\lg n)$.

---

For more details of order-statistic tree, please refer to CLRS (Section 14.1: Dynamic order statistics, 2nd edition) or this wiki article.