// subtask 2 only
#include "holiday.h"
// --- algolib/segment-tree.h ---
// Iterative segment-tree with lazy propagation, using 3n memory
// https://codeforces.com/blog/entry/18051
// Applying an Operation to a Value gets the interval size as a second parameter
// Some of the Value + Value calls might be invalid, but those results will never
// be used. This is because the tree is actually a set of trees
// TODO disable (more of) lazy propagation for Operation = NoOperation
// --- algolib/utility/bits.h ---
// will be replaced by c++20's bit header
// functions assume x >= 0
#include <bitset>
// log_floor(0) = 0
template<typename T>
int log_floor(T x) { return x ? 63 - __builtin_clzll(x) : 0; }
// log_ceil(0) = 0
template<typename T>
int log_ceil(T x) { return log_floor(x) + !!(x & (x - 1)); }
// trailing_zeros(0) = 0
template<typename T>
int trailing_zeros(T x) { return x ? __builtin_ctzll(x) : 0; }
// leading_zeros(0) = 0
template<typename T>
int leading_zeros(T x) { return x ? __builtin_clzll(x) : 0; }
template<typename T>
int popcount(T x) { return std::bitset<64>(x).count(); }
// rounds up to the next power of 2
// round_power2(0) = 0
template<typename T>
T round_power2(T x) { return x & (x - 1) ? 1 << (log_floor(x) + 1) : x; }
// ----------------
// --- algolib/utility/default-augmentations.h ---
// using this operation will (somewhat) disable lazy propagation
struct NoOperation { };
// to simplify the code, augmentation will still be done, but hopefully the compiler
// realises it's useless
// when using NoValue, NoOperation has to be used
struct NoValue {
friend NoValue operator+(const NoValue& a, const NoValue& b) { return a; }
};
// ----------------
#include <algorithm>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
template<typename Value, typename Operation>
class SegmentTree {
int n, h;
std::vector<Value> tree;
std::vector<std::optional<Operation>> lazy;
std::vector<int> roots; // only used for descend and lazily filled
// apply operation to node i and mark subtree for lazy update
void apply(int i, int k, const Operation& operation) {
if constexpr (!std::is_same_v<Operation, NoOperation>) {
tree[i](operation, k);
if (i < n) { // push to children
if (lazy[i])
lazy[i] = operation + *lazy[i];
else
lazy[i] = operation;
}
}
}
// if there's a lazy update at i, apply it and push it further down
// here k is the size of the children
void push_down(int i, int k) {
if (lazy[i]) {
apply(2*i, k, *lazy[i]);
apply(2*i + 1, k, *lazy[i]);
lazy[i].reset();
}
}
// push lazy updates from top to i >= n. bits of i encode the path
void propagate(int i) {
int k = 1 << h >> 1;
for (int a = h; a > 0; a--, k /= 2)
push_down(i >> a, k);
}
// value at i has changed. combine path to root
void combine(int i) {
i /= 2;
for (int k = 2; i > 0; i /= 2, k *= 2) {
tree[i] = tree[2*i] + tree[2*i+1];
if constexpr (!std::is_same_v<Operation, NoOperation>)
if (lazy[i]) // entire subtree affected, so just apply to i
tree[i](*lazy[i], k); // pushing it down is not necessary
}
}
public:
SegmentTree(int n_ = 1) : n(n_), h(log_floor(n)), tree(2*n),
lazy(n) { }
SegmentTree(std::vector<Value> a) : SegmentTree(a.size()) {
std::move(a.begin(), a.end(), tree.begin() + n);
for (int i = n-1; i > 0; i--)
tree[i] = tree[2*i] + tree[2*i+1];
}
Value query(int l, int r) { // [l, r[
l += n, r += n;
propagate(l), propagate(r - 1);
Value left{}, right{};
for (; l < r; l /= 2, r /= 2) {
if (l & 1)
left = left + tree[l++];
if (r & 1)
right = tree[--r] + right;
}
return left + right;
}
Value query(int i) {
i += n;
propagate(i);
return tree[i];
}
void set(int i, const Value& value) {
i += n;
propagate(i);
tree[i] = value;
combine(i);
}
void update(int l, int r, const Operation& operation) { // [l, r[
l += n, r += n;
propagate(l), propagate(r - 1);
int old_l = l, old_r = r;
for (int k = 1; l < r; l /= 2, r /= 2, k *= 2) {
if (l & 1)
apply(l++, k, operation);
if (r & 1)
apply(--r, k, operation);
}
combine(old_l), combine(old_r - 1);
}
// given some predicate left, decides at each vertex which direction to
// go, ending at some index, those (index, value) is returned.
// essentially performing a binary search on the represented data
// left(Value) = true, if the answer is contained in the subtree of Value
// index can be n ("end")
template<typename Callback>
std::pair<int, Value> descend(Callback left) {
int i = 1, k = n; // root/size in the power of 2 case
if (n & (n-1)) {
if (roots.empty()) { // fill once
std::vector<int> roots_r;
for (int l = n, r = 2*n; l < r; l /= 2, r /= 2) {
if (l & 1)
roots.push_back(l++);
if (r & 1)
roots_r.push_back(--r);
}
roots.insert(roots.end(), roots_r.rbegin(), roots_r.rend());
}
auto it = roots.begin();
while (it != roots.end() && !left(tree[*it]))
it++; // no lazy update needed, roots only apply downwards
if (it == roots.end())
return { n, Value{} };
i = *it;
k = 1;
for (int j = i; j < n; j *= 2) // size of the tree rooted at i
k *= 2;
}
if (!left(tree[i]))
return { n, Value{} };
for (k /= 2; i < n; k /= 2) { // usual descend
push_down(i, k);
i = 2*i + !left(tree[2*i]);
}
return { i - n, tree[i] };
}
};
// ----------------
#include <bits/stdc++.h>
using namespace std;
struct Sum {
long long x = 0;
int size = 0;
friend Sum operator+(const Sum& a, const Sum& b) {
return { a.x + b.x, a.size + b.size };
}
};
long long findMaxAttraction(int n, int start, int d, int attraction[]) {
vector<int> sorted(n), rank(n);
iota(sorted.begin(), sorted.end(), 0);
sort(sorted.begin(), sorted.end(), [&](int i, int j) {
return attraction[i] > attraction[j];
});
for (int i = 0; i < n; i++)
rank[sorted[i]] = i;
SegmentTree<Sum, NoOperation> tree(n);
long long ans = 0;
for (int r = 0; r < min(n, d); r++) {
tree.set(rank[r], { attraction[r], 1 });
int budget = d - r;
int lo = 0, hi = n;
while (lo+1 < hi)
if (int m = (lo + hi)/2; tree.query(0, m).size < budget)
lo = m;
else
hi = m;
ans = max(ans, tree.query(0, hi).x);
}
return ans;
}
