1928D - Lonely Mountain Dungeons - CodeForces Solution

#include <bits/stdc++.h>
#define ll long long
#define pii pair<int, int>
#define endl '\n'
#define int ll

using namespace std;

vector<int> grp;
int N, b, X;

int f(int groups) {
    if(groups == 1) return 0;
    int ans = 0;
    for(int i = 0; i < N; i++) {
        if(grp[i] == 1) continue;
        int A = grp[i] / groups + ((grp[i] % groups) > 0);
        int B = grp[i] / groups;
        int A_cnt = grp[i] % groups;
        if(A_cnt == 0) A_cnt = groups;
        int B_cnt = min(grp[i], groups) - A_cnt;

        ans += (((A_cnt - 1)*(A_cnt)*A*A)/2LL)*b;
        ans += A_cnt*B_cnt*A*B*b;
        ans += (((B_cnt - 1)*(B_cnt)*B*B)/2LL)*b;
    }

    return ans - (groups - 1)*X;
}

int ternary_search(int l, int r) {
    int prev_l = -1, prev_r = -1;
    while (l <= r) {
        if(prev_l == l && prev_r == r) {
            int mx = f(l);
            for(int i = l + 1; i <= r; i++) {
                mx = max(mx, f(i));
            }

            return mx;
        }


        prev_l = l;
        prev_r = r;
        int m1 = l + (r - l) / 3;
        int m2 = r - (r - l) / 3;
        int f1 = f(m1);
        int f2 = f(m2);
        if (f1 < f2)
            l = m1;
        else
            r = m2;
    }
    return f(l);
}

void solve() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> N >> b >> X;
    grp.resize(N);

    int mx = 0;
    for(int i = 0; i < N; i++) {
        cin >> grp[i];
        mx = max(mx, grp[i]);
    }

    cout << ternary_search(1, mx) << endl;
}

signed main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    int t;
    cin >> t;

    while(t--)
        solve();
}