Sliding window median

sol.cc

#include <map>
#include <utility>
#include <vector>
#include <algorithm>
#include <iostream>
#include <iomanip>
using namespace std;

typedef long long ll;

void show_pair(const pair<int, int> &p) {
    cout << p.first << " " << p.second << endl;
}

void show(map<int, int> &small, map<int, int> &large) {
    cout << "small\n";
    for(const auto &el : small) show_pair(el);
    cout << "large\n";
    for(const auto &el : large) show_pair(el);
}

vector<double> solve(vector<ll> &arr, ll k) {
    vector<double> ret;
    if(k == 1) {
        for(auto &el : arr) ret.push_back((double)el);
        return ret;
    } else if(k == 2) {
        for(int i = 0; i < arr.size() - 1; i++) ret.push_back((arr[i] + arr[i + 1]) / 2.0);
        return ret;
    }

    vector <ll> pre(arr.begin(), arr.begin() + k);
    sort(pre.begin(), pre.end());
    map<ll, ll> small, large;
    for (int i = 0; i < k / 2; ++i) small[pre[i]]++;
    for (int i = k / 2; i < k; ++i) large[pre[i]]++;
    ll small_size = k / 2, large_size = k - k / 2;

    ret.push_back(k % 2 ? large.begin()->first : (large.begin()->first + small.rbegin()->first) / 2.0);

    int n = arr.size();
    for(int i = k; i < n; i++) {
        // Insert the new element
        if(arr[i] <= small.rbegin()->first) {
            small[arr[i]]++;
            small_size++;
        } else {
            large[arr[i]]++;
            large_size++;
        }

        // Remove the (i - k)th element
        ll rem = arr[i - k];
        auto it = large.find(rem);
        if(it == large.end()) {
            small_size--;
            it = small.find(rem);
            it->second--;
            if(it->second == 0) small.erase(it);
        } else {
            large_size--;
            it->second--;
            if(it->second == 0) large.erase(it);
        }

        // Transition elements from small to large if small is too big.
        while(small_size > k / 2) {
            auto rit = small.rbegin();
            ll diff = small_size - k / 2;
            ll key = rit->first;
            if(rit->second > diff) {
                rit->second -= diff;
                small_size -= diff;
                large[key] += diff;
                large_size += diff;
            } else {
                diff = rit->second;
                small.erase(small.find(key));
                small_size -= diff;
                large[key] += diff;
                large_size += diff;
            }
        }

        // Transition elements from large to small if small is too small.
        while(small_size < k / 2) {
            it = large.begin();
            ll diff = k / 2 - small_size;
            if(it->second > diff) {
                it->second -= diff;
                large_size -= diff;
                small[it->first] += diff;
                small_size += diff;
            } else {
                diff = it->second;
                large.erase(it);
                large_size -= diff;
                small[it->first] += diff;
                small_size += diff;
            }
        }

        ret.push_back(k % 2 ? large.begin()->first : (large.begin()->first + small.rbegin()->first) / 2.0);
    }

    return ret;
}

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

    vector<ll> nums(n);
    for (int i = 0; i < n; ++i) cin >> nums[i];

    vector<double> sol = solve(nums, k);
    cout << fixed << setprecision(1);
    for(const auto &el : sol) cout << " " << el;
    cout << endl;
    return 0;
}