Submission #912049

#	Time	Username	Problem	Language	Result	Execution time	Memory
912049		GrindMachine	Gift (IZhO18_nicegift)	C++17	100 / 100	897 ms	85172 KiB

nicegift

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

template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;

#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl

#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)

template<typename T>
void amin(T &a, T b) {
    a = min(a,b);
}

template<typename T>
void amax(T &a, T b) {
    a = max(a,b);
}

#ifdef LOCAL
#include "debug.h"
#else
#define debug(x) 42
#endif

/*

refs:
https://oj.uz/submission/908159

sum%k != 0 or mx > sum/k => not possible
otherwise, always possible

pick the k max elems and subtract 1 from all of them
can be implemented using a priority_queue
passes the first 3 subtasks

full solution idea:
again, pick the k max elems
let their values be v1 <= v2 <= ... <= vk
we can subtract at most v1 from all these guys
because the maximum we can subtract is v1, lets try to subtract v1
when will this be invalid?

recollect the condition from the start:
sum%k == 0 and mx <= sum/k for a solution to exist
the value of sum%k doesnt change after any operation, so if it's 0 in the beginning, then we dont have to care about it
what about mx <= sum/k?
lets say we subtract x

1) new maximum = vk
we want vk-x <= (sum-k*x)/k
after simplification, we get vk <= sum/k, which is already satisfied
so we dont have to worry if the new maximum = vk

2) new maximum = max guy that is not chosen in the op (let it be w)
w <= (sum-k*x)/k
after simplification, we get x <= (sum-w*k)/k
we cant subtract v1 if w > (sum-v1*k)/k
in this case, the max value we can subtract is (sum-w*k)/k
otherwise, we can subtract v1

x = min(v1,(sum-w*k)/k)

if v1 is subtracted, then at least 1 guy becomes 0 in the op
if (sum-w*k)/k is subtracted, then no one may become 0 in this op
but it's guaranteed that at least 1 guy becomes 0 in the next op (because after the op, w = sum/k)
so the #of ops would be minimal

gives a bound of 2*n*k guys in total
this is a loose bound, probably possible to prove a tighter bound because the solution expects around 1.5*n*k ops, but not sure how to do it

*/

const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

void solve(int test_case)
{
    ll n,k; cin >> n >> k;
    vector<ll> a(n+5);
    rep1(i,n) cin >> a[i];
 
    ll sum = 0;
    rep1(i,n) sum += a[i];
 
    if(sum%k){
        cout << -1 << endl;
        return;
    }
 
    ll mx = *max_element(all(a));
    if(mx > sum/k){
        cout << -1 << endl;
        return;
    }
 
    priority_queue<pll> pq;
    rep1(i,n) pq.push({a[i],i});
 
    vector<vector<ll>> ans;
 
    while(true){        
        if(!pq.top().ff) break;

        vector<ll> curr;
        curr.pb(0);
        
        rep1(iter,k){
            auto [val,i] = pq.top();
            pq.pop();
            curr.pb(i);
        }

        ll curr_mx = a[curr[1]];
        ll curr_mn = a[curr.back()];

        ll remain_mx = -inf2;
        if(!pq.empty()){
            remain_mx = pq.top().ff;
        }

        // what is the max val that can be subtracted?
        
        /*

        ll l = 1, r = curr_mn;
        ll mx_sub = -1;

        while(l <= r){
            ll mid = (l+r) >> 1;
            ll new_sum = sum-mid*k;
            if(curr_mx-mid <= new_sum/k and remain_mx <= new_sum/k){
                mx_sub = mid;
                l = mid+1;
            }
            else{
                r = mid-1;
            }
        }

        assert(mx_sub != -1);
    
        */

        ll mx_sub = min(curr_mn,(sum-remain_mx*k)/k);

        curr[0] = mx_sub;
        sum -= mx_sub*k;

        rep1(ind,sz(curr)-1){
            ll i = curr[ind];
            a[i] -= mx_sub;
            pq.push({a[i],i});
        }

        ans.pb(curr);
    }
    
    assert(sum == 0);
    rep1(i,n){
        assert(a[i] == 0);
    }
 
    cout << sz(ans) << endl;
    trav(ar,ans){
        trav(i,ar) cout << i << " ";
        cout << endl;
    }
}

int main()
{
    fastio;

    int t = 1;
    // cin >> t;

    rep1(i, t) {
        solve(i);
    }

    return 0;
}

Compilation message (stderr)

nicegift.cpp: In function 'void solve(int)':
nicegift.cpp:140:12: warning: unused variable 'curr_mx' [-Wunused-variable]
  140 |         ll curr_mx = a[curr[1]];
      |            ^~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 11 / 11
• Subtask 3: 12 / 12
• Subtask 4: 19 / 19
• Subtask 5: 51 / 51