Explanation

Each query asks us to remove the element at a given position from what currently remains in the list.

The naive approach would be to shift the elements in the array after each removal, which costs O(n)O(n) per query, which doesn't suffice here given the constraints. We need a data structure that supports both k-th element queries and deletions efficiently.

All three solutions below process each query in O(logn)O(\log n) or O(log2n)O(\log^2 n), giving O(nlogn)O(n \log n) or O(nlog2n)O(n\log^2 n) overall.

Solution 1 - Indexed Set

Store the original indices of all still-present elements in an indexed set. This C++ builtin data structure can process queries of the form "remove element at position pp", and can also access the original index at a position in O(logn)O(\log n) time using find_by_order.

#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
template <class T>
using OrderedSet =
tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

Solution 2 - Binary Search on BIT

We maintain a Binary Indexed Tree (BIT) where each position is initialized to 1 (present) and gets set to 0 when removed. The prefix sum prefix_sum(i) then counts how many elements are still present in positions 11 through ii.

To find the original index of the kk-th remaining element, binary search for the smallest index where prefix_sum(index) == k. Each step of the binary search does a BIT query, giving O(log2n)O(\log^2 n) per operation.

#include <bits/stdc++.h>
using namespace std;
int n;
Code Snippet: BIT Code (from PURS module) (Click to expand)
// Binary search for smallest index where prefix_sum == target
int find_kth(int target, BIT<int> &bit) {
int lo = 1, hi = n;

Solution 3 - Segment Tree

Build a segment tree where each node stores the count of the present elements in its range. To find the kk-th remaining element, walk down the tree as follows: if the left child contains k\geq k elements, recurse left; otherwise, subtract the left child's count from kk and recurse right.

This finds the answer in O(logn)O(\log n) by descending exactly one root-to-leaf path.

#include <bits/stdc++.h>
using namespace std;
Code Snippet: Segment Tree (Click to expand)
int main() {
int n;
cin >> n;

Join the USACO Forum!

Stuck on a problem, or don't understand a module? Join the USACO Forum and get help from other competitive programmers!