PROBLEM LINK:
Author: Roman Furko
Testers: Pushkar Mishra and Sergey Kulik
Editorialist: Pawel Kacprzak
Russian Translator: Sergey Kulik
Mandarian Translator: Gedi Zheng
DIFFICULTY:
Medium
PREREQUISITES:
balanced binary serach data structures with fast union (treap, AVL-tree, splay tree), order statistics
PROBLEM:
You are given N sets of integers. Initially each set consists of exactly one element. All elements are unique. Your task is to be able to perform two operations on these sets:
- Merge two sets into a new one
- For a given set, return the k^{th}
smallest integer in it
QUICK EXPLANATION:
Use any balanced binary search tree with the ability to perform fast unions to represent each set. Augment each tree in order to provide a method to return the k^{th} smallest element in it.
EXPLANATION:
We will show how to use a treap as the underlaying data structure.
Finding the k^{th} element in a treap
In every node of the treap, store the number of nodes in its left subtree. Notice that this value can be easily updated and maintained during any rotation used in treap implementation. Having this information, we can easily return the k^{th} smallest element in the treap. Starting our recursive procedure in with the whole treap as a current search space, let c be the number of nodes in root’s left subtree. We have 3 cases to consider:
-
c = k - 1, then element in the
root is the element we search for,
because in the current search space. -
c >= k, then we know that the
element we are searching for is in
the left subtree of the root -
c < k - 1, then we search for the
(k - c - 1)^{th} element in the right subtree
of the root
This allows us to perform a single query in O(\log(M)), where M is the number of nodes in the treap.
Handling union of two treaps
Using the fact that all our input numbers are unique, we can easily use a recursive algorithm to find the union of two treaps.
These operations were described widely in the past, so I encourage all of you to read about them, for example, here.
A single union takes O(M \cdot \log(\frac{N}{M})), where M and N are the sizes of two treaps and M \leq N.
Time complexity
The exact complexity depends on the operations given in a test file, because different operations provide different treaps sizes, but notice that the time needed for union of two heaps is dominated by the size of the smaller treap. In order to produce two treaps with size \frac{N}{c}, you have to perform 2 \cdot \frac{N}{c} - 2 unions beforehand. In summary, these method will easily pass all testcases.