Sliding Cost

You are given an array of n integers. Your task is to calculate for each window of k elements, from left to right, the minimum total cost of making all elements equal.

You can increase or decrease each element with cost x where x is the difference between the new and the original value. The total cost is the sum of such costs.

Example:

Input:  n = 8, k = 3, a = {2, 4, 3, 5, 8, 1, 2, 1}
Output: 2 2 5 7 7 1 

Approach

C++

#include <bits/stdc++.h>
using namespace std;

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
template <class T>
using ordered_set = tree<T, null_type,
                         less<T>, rb_tree_tag,
                         tree_order_statistics_node_update>;
template <class K, class V>
using ordered_map = tree<K, V, less<K>,
                         rb_tree_tag,
                         tree_order_statistics_node_update>;

vector<long long> slidingCost(int n, int k, vector<int> &a)
{

    ordered_set<array<int, 2>> os;
    for (int i = 0; i < k; i++)
        os.insert({a[i], i});
    int med = (*os.find_by_order((k - 1) / 2))[0];
    long long ans = 0;
    for (int i = 0; i < k; i++)
        ans += abs(a[i] - med);
    vector<long long> res;
    res.push_back(ans);
    for (int i = 0, j = k; j < n; i++, j++)
    {
        os.erase({a[i], i});
        os.insert({a[j], j});
        int pre = med;
        med = (*os.find_by_order((k - 1) / 2))[0];
        ans -= abs(a[i] - pre);
        ans += abs(a[j] - med);
        if (k % 2 == 0)
            ans -= (med - pre);
        res.push_back(ans);
    }
    return res;
}

int main()
{
    int n = 8, k = 3;

    vector<int> a = {2, 4, 3, 5, 8, 1, 2, 1};

    vector<long long> res = slidingCost(n, k, a);

    for (int i = 0; i < res.size(); i++)
        cout << res[i] << " ";

    return 0;
}


No comments:

Post a Comment