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;
}

Please help me, I don't know what is wrong. Ploblem Link

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
en1 English achhadahappy 2023-06-01 16:30:27 4570 Initial revision (published)