#include "ricehub.h"
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int besthub(int RiceCount, int Length, int RiceFields[], long long Budget)
{
if (RiceCount <= 5000) {
ll best_count = 0;
for (ll i = 0; i < RiceCount; i++)
{
ll x = RiceFields[i];
ll cost = 0;
ll count = 1;
ll backwards_index = i-1;
ll forwards_index = i+1;
while (cost <= Budget) {
if (backwards_index < 0 && forwards_index >= RiceCount) {
return RiceCount;
}
ll backwards = backwards_index < 0 ? 1e9 : x-RiceFields[backwards_index];
ll forwards = forwards_index >= RiceCount ? 1e9 : RiceFields[forwards_index]-x;
if (forwards < backwards) {
cost += forwards;
forwards_index++;
count++;
} else {
cost += backwards;
backwards_index--;
count++;
}
}
count --;
best_count = max(best_count, count);
}
return best_count;
}
ll count_before = 0;
ll count_after = 0;
ll cost = 0;
ll best_count = 1;
for (ll i = 1; i < RiceCount; i++)
{
// set the hub at this field
ll x = RiceFields[i];
ll distance = RiceFields[i] - RiceFields[i-1];
// account for the prev_index now also having distance
count_before++;
cost += distance * count_before;
cost -= distance * count_after;
if (count_after > 0) {
count_after --;
}
auto add_after = [&](){
while (cost <= Budget && i+count_after < RiceCount-1) {
count_after++;
cost += RiceFields[i+count_after] - x;
}
if (cost > Budget && count_after > 0) {
cost -= RiceFields[i+count_after] - x;
count_after--;
}
};
add_after();
while (
count_before > 0
&& (
cost > Budget
|| (
i+count_after < RiceCount-1 &&
(x - RiceFields[i-count_before]) >= (RiceFields[i+count_after+1] - x)
)
)) {
cost -= x - RiceFields[i-count_before];
count_before--;
add_after();
}
ll current_count = count_before + count_after + 1;
if (current_count > best_count) {
best_count = current_count;
}
}
return best_count;
}
