Advertisement
coloriot

HA_42_3

Feb 16th, 2025
28
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.60 KB | None | 0 0
  1. #include <iostream>
  2. #include <vector>
  3. #include <algorithm>
  4. #include <cassert>
  5. using namespace std;
  6.  
  7. typedef long long ll;
  8.  
  9. inline ll Lcost(int p, int r, const vector<ll>& b, const vector<ll>& prefix) {
  10.     if(r == 0) return 0;
  11.     ll sumLeft = prefix[p-1] - ((p - r - 1) >= 0 ? prefix[p - r - 1] : 0LL);
  12.     return (ll)r * b[p] - sumLeft;
  13. }
  14.  
  15. inline ll Rcost(int p, int s, const vector<ll>& b, const vector<ll>& prefix) {
  16.     if(s == 0) return 0;
  17.     return (prefix[p+s] - prefix[p]) - (ll)s * b[p];
  18. }
  19.  
  20. inline ll costForPivot(int p, int r, int k, const vector<ll>& b, const vector<ll>& prefix, int n) {
  21.     int s = k - r;
  22.     return Lcost(p, r, b, prefix) + Rcost(p, s, b, prefix);
  23. }
  24.  
  25. int findBestR(int p, int lo, int hi, int k, const vector<ll>& b, const vector<ll>& prefix, int n) {
  26.     while(lo < hi) {
  27.         int mid = (lo + hi) / 2;
  28.         ll f_mid = costForPivot(p, mid, k, b, prefix, n);
  29.         ll f_next = costForPivot(p, mid+1, k, b, prefix, n);
  30.         if(f_mid > f_next)
  31.             lo = mid + 1;
  32.         else
  33.             hi = mid;
  34.     }
  35.     return lo;
  36. }
  37.  
  38. class NearestNeighborsSolver {
  39. public:
  40.     void solve() {
  41.         ios::sync_with_stdio(false);
  42.         cin.tie(nullptr);
  43.  
  44.         int n, k;
  45.         cin >> n >> k;
  46.         vector<ll> a(n);
  47.         for (int i = 0; i < n; i++){
  48.             cin >> a[i];
  49.         }
  50.  
  51.         vector<pair<ll,int>> arr(n);
  52.         for (int i = 0; i < n; i++){
  53.             arr[i] = {a[i], i};
  54.         }
  55.         sort(arr.begin(), arr.end(), [](auto &p1, auto &p2){
  56.             return p1.first < p2.first;
  57.         });
  58.  
  59.         vector<ll> b(n);
  60.         vector<int> sortedToOrig(n);
  61.         for (int i = 0; i < n; i++){
  62.             b[i] = arr[i].first;
  63.             sortedToOrig[i] = arr[i].second;
  64.         }
  65.  
  66.         vector<ll> prefix(n, 0);
  67.         prefix[0] = b[0];
  68.         for (int i = 1; i < n; i++){
  69.             prefix[i] = prefix[i-1] + b[i];
  70.         }
  71.  
  72.         vector<ll> ans(n, 0);
  73.         for (int p = 0; p < n; p++){
  74.             int r_low = max(0, k - ((n - 1) - p));
  75.             int r_high = min(k, p);
  76.             int bestR = r_low;
  77.             if(r_low < r_high)
  78.                 bestR = findBestR(p, r_low, r_high, k, b, prefix, n);
  79.             ll bestCost = costForPivot(p, bestR, k, b, prefix, n);
  80.             ans[ sortedToOrig[p] ] = bestCost;
  81.         }
  82.  
  83.         for (int i = 0; i < n; i++){
  84.             cout << ans[i] << (i+1==n ? "\n" : " ");
  85.         }
  86.     }
  87. };
  88.  
  89. int main(){
  90.     ios::sync_with_stdio(false);
  91.     cin.tie(nullptr);
  92.  
  93.     NearestNeighborsSolver solver;
  94.     solver.solve();
  95.     return 0;
  96. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement