1077 - Sliding Cost

Click to show code.


using namespace std;
using ll = long long;
using ii = pair<ll, ll>;
using vi = vector<ll>;
template <typename InputIterator,
          typename T = typename iterator_traits<InputIterator>::value_type>
void write(InputIterator first, InputIterator last, const char *delim = "\n")
{
    copy(first, last, ostream_iterator<T>(cout, delim));
}
template <typename InputIterator,
          typename T = typename iterator_traits<InputIterator>::value_type>
void read_n(InputIterator it, int n)
{
    copy_n(istream_iterator<T>(cin), n, it);
}
ii med = {-1, -1};
set<ii> bot, top;
ll sbot, stop;
void fix()
{
    int m = 1 + (int)bot.size() + (int)top.size();
    if ((int)bot.size() < (m - 1) / 2)
    {
        bot.insert(med), sbot += med.first;
        med = *top.begin();
        top.erase(med), stop -= med.first;
    }
    if ((int)bot.size() > (m - 1) / 2)
    {
        top.insert(med), stop += med.first;
        med = *prev(bot.end());
        bot.erase(med), sbot -= med.first;
    }
}
void add(ii x)
{
    if (med.first == -1)
    {
        med = x;
        return;
    }
    if (x < med)
        bot.insert(x), sbot += x.first;
    else
        top.insert(x), stop += x.first;
    fix();
}
void rem(ii x)
{
    if (x == med)
    {
        med = *top.begin();
        top.erase(med), stop -= med.first;
    }
    else if (x < med)
        bot.erase(x), sbot -= x.first;
    else
        top.erase(x), stop -= x.first;
    fix();
}
int main(void)
{
    ios::sync_with_stdio(false), cin.tie(NULL);
    int n, k;
    cin >> n >> k;
    vi a(n);
    read_n(a.begin(), n);
    if (k == 1)
    {
        for (int i = 0; i < n; ++i)
            cout << 0 << " ";
        cout << endl;
        return 0;
    }
    for (int i = 0; i < n; ++i)
    {
        add({a[i], i});
        if (i >= k - 1)
        {
            auto left = ((int)bot.size() * med.first - sbot);
            auto right = (stop - (int)top.size() * med.first);
            cout << left + right << " ";
            rem({a[i - k + 1], i - k + 1});
        }
    }
    return 0;
}