#include <algorithm>
#include <functional>
#include <vector>
template <typename T> class SegmentTree {
public:
int n;
T idt;
std::vector<T> seg;
std::function<T(T, T)> f;
SegmentTree(int n, std::function<T(T, T)> f = std::plus<T>(), T idt = T()) : n(n), idt(idt), f(f), seg(2 * n, idt) {}
void set(int idx, T x) {
for (seg[idx += n] = x, idx /= 2; idx > 0; idx /= 2) {
seg[idx] = f(seg[2 * idx], seg[2 * idx + 1]);
}
}
T query(int l, int r) {
T ansL = idt, ansR = idt;
for (l += n, r += n + 1; l < r; l /= 2, r /= 2) {
if (l & 1) ansL = f(ansL, seg[l++]);
if (r & 1) ansR = f(seg[--r], ansR);
}
return f(ansL, ansR);
}
int partition_point(int l, const std::function<bool(T)> &t) {
T p = idt;
for (l += n; t(f(p, seg[l])) and l; l /= 2) {
if (l & 1 and l != 1) p = f(p, seg[l++]);
}
if (!l) {
return n;
}
while (l < n) {
if (t(f(p, seg[l <<= 1]))) p = f(p, seg[l++]);
}
return l - n;
}
};
long long int solve(int n, int start, int d, std::vector<int> attraction) {
auto build_a = [&]() {
std::vector<std::pair<int, int>> a;
for (int i = 0; i < n; ++i) {
a.push_back({attraction[i], i});
}
std::sort(a.rbegin(), a.rend());
std::vector<std::pair<int, int>> new_a(n);
for (int i = 0; i < n; ++i) {
new_a[a[i].second] = {a[i].first, i};
}
return a = std::move(new_a);
};
std::vector<std::pair<int, int>> a = build_a();
struct Node {
int count = 0;
int64_t sum = 0;
};
auto round_up_pow2 = [](int n) {
--n;
n |= n >> 1;
n |= n >> 2;
n |= n >> 4;
n |= n >> 8;
n |= n >> 16;
return n + 1;
};
SegmentTree<Node> st(round_up_pow2(n), [&](Node a, Node b) {
return Node({a.count + b.count, a.sum + b.sum});
});
auto calc_f = [&](auto &&self, int tl, int tr, int l, int r, int start, std::vector<std::pair<int, int64_t>> &f) {
int tm = (tl + tr) / 2;
std::pair<int64_t, int> best = {-1, -1};
for (int i = l; i <= r; ++i) {
st.set(a[i].second, {1, a[i].first});
int idx = st.partition_point(0, [&](Node x) {
return x.count <= tm - (i - start);
}) - 1;
if (idx == -1) {
continue;
}
best = std::max(best, {st.query(0, idx).sum, -i});
}
best.second = -best.second;
f[tm] = {best.second, best.first};
if (tl == tr) {
return;
}
for (int i = f[tm].first; i <= r; ++i) {
st.set(a[i].second, Node{});
}
self(self, tm + 1, tr, f[tm].first, r, start, f);
for (int i = l; i <= r; ++i) {
st.set(a[i].second, Node{});
}
self(self, tl, tm, l, f[tm].first, start, f);
};
std::vector<std::pair<int, int64_t>> f1(d + 1);
calc_f(calc_f, 1, d, start, n - 1, start, f1);
for (int i = 0; i < n; ++i) {
st.set(i, Node{});
}
std::reverse(attraction.begin(), attraction.begin() + start);
for (int i = start; i < n; ++i) {
attraction[i] = 0;
}
a = build_a();
std::vector<std::pair<int, int64_t>> f2(d + 1);
calc_f(calc_f, 1, d, 0, n - 1, 0, f2);
long long int ans = 0;
for (int i = 1; i <= d; ++i) {
long long int cur = f1[i].second;
int rem = d - i - f1[i].first + start - 1;
if (rem > 0) {
cur += f2[rem].second;
}
ans = std::max(ans, cur);
}
return ans;
}
long long int findMaxAttraction(int n, int start, int d, int attraction[]) {
std::vector<int> a(n);
for (int i = 0; i < n; ++i) {
a[i] = attraction[i];
}
long long int ans1 = solve(n, start, d, a);
std::reverse(a.begin(), a.end());
start = n - start - 1;
long long int ans2 = solve(n, start, d, a);
return std::max(ans1, ans2);
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |