Please subscribe to the official Codeforces channel in Telegram via the link https://t.me/codeforces_official. ×

Needed help for Codechef Problem Beautification of Array

Revision en1, by achhadahappy, 2023-06-01 16:30:27
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define ll long long
#define int long long
#define endl "\n"
#define ptr(a, n)               \
for (int i = 0; i < n; i++) \
{                           \
cout << a[i] << ' ';    \
}
#define fort(z, x, t) for (ll z = x; z < t; z++)
#define rev(z, x, t) for (ll z = t; z > x; z--)
#define arin(a, n)             \
ll a[n];                   \
for (ll i = 0; i < n; i++) \
{                          \
cin >> a[i];           \
}
#define pb push_back
#define vein(v, n)             \
vector<ll> v;              \
for (ll i = 0; i < n; i++) \
{                          \
ll x;                  \
cin >> x;              \
v.pb(x);               \
}
#define seint(s, n)            \
set<int> s;                \
for (ll i = 0; i < n; i++) \
{                          \
ll x;                  \
cin >> x;              \
s.insert(x);           \
}
#define vint(v) vector<ll> v;
#define vstring(vs, n)         \
vector<string> vs;         \
for (ll i = 0; i < n; i++) \
{                          \
string s;              \
cin >> s;              \
vs.pb(s);              \
}
#define num(n) \
ll n;      \
cin >> n;
#define out(n) cout << n << endl;
#define sortv(v) sort(v.begin(), v.end());
#define mpll map<ll, ll> mpi;
#define mpchar map<char, ll> mpc;
#define cin(a, n, b) \
f(i, n)          \
{                \
cin >> a[i]; \
b[i] = a[i]; \
}
#define yes cout << "YES" << endl;
#define no cout << "NO" << endl;
typedef tree<ll, null_type, less_equal<ll>, rb_tree_tag, tree_order_statistics_node_update> ordered_multiset;
typedef tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;
const int N = 1e5 + 10;
int mod = 1e9 + 7;
int fun(int mid, int k, vector<int> &v)
{
int n = v.size();
int val = 0;
for (int i = 0; i < n; i++)
{
val += max(0ll, mid - v[i]);
if (val > k)
return 0;
}
if (val > k)
return 0;
return 1;
}
void solve()
{
num(n);
num(k);
vein(v, n);
sortv(v);
int req = 0;
for (int i = 0; i < n; i++)
{
req += max(0ll, v.back() - v[i]);
}
if (req <= k)
{
k -= req;
int val1 = k / n;
int val2 = k % n;
for (int i = 0; i < n; i++)
{
v[i] = val1 + v.back();
}
sortv(v);
for (int i = 0; i < val2; i++)
v[i]++;
vector<int> suf(n + 1, 0);
suf[n - 1] = v[n - 1];
for (int i = n - 2; i >= 0; i--)
{
suf[i] = suf[i + 1] + v[i];
}
int ans = 0;
for (int i = 0; i < n - 1; i++)
{
ans += (v[i] * suf[i + 1]) % mod;
ans %= mod;
}
out(ans % mod);
return;
}
sortv(v);
int miin = v[0];
int maax = v.back();
while (miin + 1 < maax)
{
int mid = (miin + maax) / 2;
if (fun(mid, k, v))
{
miin = mid;
}
else
{
maax = mid;
}
}
for (int i = 0; i < n; i++)
{
if (v[i] < miin)
{
int val2 = min(k, miin - v[i]);
v[i] += val2;
k -= val2;
}
}
sortv(v);
for (int i = 0; i < k; i++)
{
v[i]++;
}
sortv(v);
vector<int> suf(n + 1, 0);
suf[n - 1] = v[n - 1];
for (int i = n - 2; i >= 0; i--)
{
suf[i] = suf[i + 1] + v[i];
}
int ans = 0;
for (int i = 0; i < n - 1; i++)
{
ans += (v[i] * suf[i + 1]) % mod;
ans %= mod;
}
// ptr(v, n);
// cout << endl;
out(ans % mod);
}
signed main()
{
// #ifndef ONLINE_JUDGE
// freopen("input.txt", "r", stdin);
// freopen("output.txt", "w", stdout);
// #endif
std::ios::sync_with_stdio(false);
std::cin.tie(0);
int t;
cin >> t;
fort(i, 0, t)
solve();
return 0;
}