#include "holiday.h"
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
struct SegmentTree {
int n, h;
vector<ll> base;
vector<int> cnt;
vector<ll> sum;
SegmentTree() {}
SegmentTree(int _n, const vector<ll>& arr) : n(_n) {
h = Log2(n);
n = 1 << h;
base.resize(2*n, 0);
for (int i=0; i<arr.size(); ++i) base[n+i] = arr[i];
cnt.resize(2*n, 0);
sum.resize(2*n, 0);
}
void update(int p, int v) {
p += n;
cnt[p] += v;
sum[p] += base[p] * v;
for (p/=2; p>0; p/=2) {
cnt[p] = cnt[2*p] + cnt[2*p+1];
sum[p] = sum[2*p] + sum[2*p+1];
}
}
ll query(int k) {
if (cnt[1] <= k) return sum[1];
ll ret = 0;
int curr = 1;
while (curr < n) {
if (cnt[2*curr+1] <= k) {
ret += sum[2*curr+1];
k -= cnt[2*curr+1];
curr = 2*curr;
} else {
curr = 2*curr+1;
}
}
return ret;
}
static int Log2(int x){
int ret = 0;
while (x > (1 << ret)) ret++;
return ret;
}
};
static const int MAX_N = 1e5 + 1;
static ll ans = 0;
static int n, start, d;
static int arr[MAX_N];
static SegmentTree seg;
static pii curr_range;
static void move(int l, int r) {
for (int& x = curr_range.first; x < l; ++x) seg.update(arr[x], -1);
for (int& x = curr_range.first; x > l; --x) seg.update(arr[x-1], 1);
for (int& x = curr_range.second; x < r; ++x) seg.update(arr[x+1], 1);
for (int& x = curr_range.second; x > r; --x) seg.update(arr[x], -1);
assert(curr_range.first == l && curr_range.second == r);
}
static void test() {
for (int i=0; i<n; ++i) cout << seg.query(i) << ' ';
cout << endl;
}
static void solve(int ls, int le, int rs, int re) {
if (ls > le) return;
assert(rs <= re);
pii past_range = curr_range;
int m = (ls + le) / 2;
ll opt_val = -1;
int opt_idx = -1;
// cout << "visit: " << ls << ' ' << le << ' ' << rs << ' ' << re << endl;
for (int i=rs; i<=re; ++i) {
// test();
move(m, i);
ll tmp = seg.query(d-2*(start-m)-(i-start));
if (opt_val < tmp) {
opt_val = tmp;
opt_idx = i;
}
}
ans = max(ans, opt_val);
// cout << m << ' ' << opt_val << ' ' << opt_idx << endl;
solve(ls, m-1, rs, opt_idx);
solve(m+1, le, opt_idx, re);
move(past_range.first, past_range.second);
}
ll findMaxAttraction(int N, int Start, int D, int attraction[]) {
n = N;
start = Start;
d = D;
vector<ll> attract(n);
for (int i=0; i<n; ++i) attract[i] = attraction[i];
for (int i=0; i<n; ++i) arr[i] = attract[i];
sort(attract.begin(), attract.end());
seg = SegmentTree(n, attract);
for (int i=0; i<n; ++i) arr[i] = lower_bound(attract.begin(), attract.end(), arr[i]) - attract.begin();
seg.update(arr[start], 1);
curr_range = { start, start };
// solve(0, start, start, n-1);
for (int i=0; i<n; ++i) {
move(0, i);
ans = max(ans, seg.query(d-i));
}
// reverse(arr, arr+n);
// start = n-1-start;
// curr_range = { start, start };
// solve(0, start, start, n-1);
return ans;
}
