Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <iostream>
- #include <vector>
- #include <algorithm>
- #include <cassert>
- using namespace std;
- typedef long long ll;
- inline ll Lcost(int p, int r, const vector<ll>& b, const vector<ll>& prefix) {
- if(r == 0) return 0;
- ll sumLeft = prefix[p-1] - ((p - r - 1) >= 0 ? prefix[p - r - 1] : 0LL);
- return (ll)r * b[p] - sumLeft;
- }
- inline ll Rcost(int p, int s, const vector<ll>& b, const vector<ll>& prefix) {
- if(s == 0) return 0;
- return (prefix[p+s] - prefix[p]) - (ll)s * b[p];
- }
- inline ll costForPivot(int p, int r, int k, const vector<ll>& b, const vector<ll>& prefix, int n) {
- int s = k - r;
- return Lcost(p, r, b, prefix) + Rcost(p, s, b, prefix);
- }
- int findBestR(int p, int lo, int hi, int k, const vector<ll>& b, const vector<ll>& prefix, int n) {
- while(lo < hi) {
- int mid = (lo + hi) / 2;
- ll f_mid = costForPivot(p, mid, k, b, prefix, n);
- ll f_next = costForPivot(p, mid+1, k, b, prefix, n);
- if(f_mid > f_next)
- lo = mid + 1;
- else
- hi = mid;
- }
- return lo;
- }
- class NearestNeighborsSolver {
- public:
- void solve() {
- ios::sync_with_stdio(false);
- cin.tie(nullptr);
- int n, k;
- cin >> n >> k;
- vector<ll> a(n);
- for (int i = 0; i < n; i++){
- cin >> a[i];
- }
- vector<pair<ll,int>> arr(n);
- for (int i = 0; i < n; i++){
- arr[i] = {a[i], i};
- }
- sort(arr.begin(), arr.end(), [](auto &p1, auto &p2){
- return p1.first < p2.first;
- });
- vector<ll> b(n);
- vector<int> sortedToOrig(n);
- for (int i = 0; i < n; i++){
- b[i] = arr[i].first;
- sortedToOrig[i] = arr[i].second;
- }
- vector<ll> prefix(n, 0);
- prefix[0] = b[0];
- for (int i = 1; i < n; i++){
- prefix[i] = prefix[i-1] + b[i];
- }
- vector<ll> ans(n, 0);
- for (int p = 0; p < n; p++){
- int r_low = max(0, k - ((n - 1) - p));
- int r_high = min(k, p);
- int bestR = r_low;
- if(r_low < r_high)
- bestR = findBestR(p, r_low, r_high, k, b, prefix, n);
- ll bestCost = costForPivot(p, bestR, k, b, prefix, n);
- ans[ sortedToOrig[p] ] = bestCost;
- }
- for (int i = 0; i < n; i++){
- cout << ans[i] << (i+1==n ? "\n" : " ");
- }
- }
- };
- int main(){
- ios::sync_with_stdio(false);
- cin.tie(nullptr);
- NearestNeighborsSolver solver;
- solver.solve();
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement