Ordered Sets

ordered_set is a regular set but with two extra functionalities:

find_by_order(k): returns an iterator for the k-th smallest element (0-indexed) in \(O(\log n)\).
order_of_key(x): returns how many elements are strictly smaller than x in \(O(\log n)\).

Implementation

To be able to make an ordered set we need to write the following globally:

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

template<typename T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;

Now we’re ready to use it. To make an ordered set we simply write ordered_set<T>, where T is the desired type (e.g, int, long long, string, etc).

Sliding Window Median

CSES easy

Given array \(x\) of \(n\) elements and a value \(k\). Print the median of all subarrays of length \(k\).

The median is the middle element when the elements are sorted. If the number of elements is even, there are two possible medians and we assume that the median is the smaller of them.

Solution

We can solve this task using ordered_set by applying a sliding window over all subarrays of length \(k\). For each window, we find the \(\frac{k}{2}\)-th smallest element using find_by_order(k/2) (note that if \(k\) is even, we instead take k/2 - 1).

Like set, ordered_set does not allow duplicates. To handle this, we represent the original array as pairs, where the first element is the value and the second is its index, ensuring that all elements are distinct.

Since we perform one insertion and one deletion for each shift of the window, the overall time complexity is \(O(n \log n)\).

#include <iostream>
using namespace std;

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

template<typename T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

int main() {
    int n, k;
    cin >> n >> k;

    pair<int, int> a[n];
    for (int i = 0; i < n; i++) {
        cin >> a[i].first;
        a[i].second  = i;
    }

    ordered_set<pair<int ,int> > os;

    //enter the first window
    for (int i = 0 ; i < k; i++){
        os.insert(a[i]);
    }

    //print the median for the first window
    int median;
    if (k % 2 == 0){
        median = k/2 - 1;
    }
    else{
        median = k/2;
    }

    pair<int ,int> ans = *os.find_by_order(median);
    cout << ans.first << ' ';

    // move through all windows
    for( int i = 0 ; i < n - k ; i++){
        //go to the next window
        //removing the first element
        os.erase(a[i]);

        // adding the one after last
        os.insert(a[i + k]);

        //print the solution of this window
        pair<int , int> ans = *os.find_by_order(median);
        cout << ans.first << ' ';
    }

}