-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathsliding-cost.cpp
78 lines (77 loc) · 1.67 KB
/
sliding-cost.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
#include <bits/stdc++.h>
using namespace std;
int main() {
int n, k;
cin >> n >> k;
// HACK: Allocate more space to avoid going out of bounds.
vector<int> a(n + 1);
for (int i = 0; i < n; i++) {
cin >> a[i];
}
set<pair<long long, int>> ls, rs;
long long L = 0, R = 0;
for (int i = 0; i < k; i++) {
ls.insert({a[i], i});
L += a[i];
}
// rebalance (ls -> rs)
while (ls.size() > k / 2 + k % 2) {
auto lmx = *ls.rbegin();
rs.insert(lmx);
R += lmx.first;
ls.erase(lmx);
L -= lmx.first;
}
auto cost = [&](long long median) {
return (median * ls.size() - L) + (R - median * rs.size());
};
int i = 0, j = k;
while (j <= n) {
auto m = *ls.rbegin();
long long c = cost(m.first);
// if k is even - try both medians
if (k % 2 == 0) {
auto ma = *rs.begin();
long long ca = cost(ma.first);
c = min(c, ca);
}
cout << c << " ";
// insert a[j]
pair<long long, int> aj = {a[j], j};
if (aj > m) {
rs.insert(aj);
R += a[j];
} else {
ls.insert(aj);
L += a[j];
}
// erase a[i]
pair<long long, int> ai = {a[i], i};
if (ls.count(ai) > 0) {
ls.erase(ai);
L -= a[i];
} else {
rs.erase(ai);
R -= a[i];
}
// rebalance (ls -> rs)
while (ls.size() > k / 2 + k % 2) {
auto lmx = *ls.rbegin();
rs.insert(lmx);
R += lmx.first;
ls.erase(lmx);
L -= lmx.first;
}
// rebalance (rs -> ls)
while (ls.size() < k / 2 + k % 2) {
auto rmn = *rs.begin();
ls.insert(rmn);
L += rmn.first;
rs.erase(rmn);
R -= rmn.first;
}
i++, j++;
}
cout << "\n";
return 0;
}