1622C - Set or Decrease

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int main(){
    ios::sync_with_stdio(0);cin.tie(0);
    ll t; cin>>t;
    for(ll i=0;i<t;i++){
        ll n; cin>>n;
        ll k; cin>>k;
        ll a[n];
        ll steps=0;
        ll sum=0; // sum of integer array "a"
        ll d=0; // d is the number of times the minimum element is decreased
        ll next=n-1; // u is the next biggest element to set
        
        for(ll j=0;j<n;j++){
            cin>>a[j];
            sum+=a[j];
        }
        sort(a,a+n); // sort the array in increasing order
        
        // Problematic code: O(a_i) time complexity
        // This while loop may take up to ~10^9 loops since a_i <= 10^9
        while(sum>k){
            // Greedy algorithm: Choose to set or decrease
            // based on which decreases the sum more
            // It's optimal for "decrease"s to be done before the "set"s
            if(next>0 and a[next]-a[0]+d>n-next){
                // choose a_i to be the largest integer that hasn't yet been set
                // choose a_j to be the smallest integer in the array
                // set a_i equal to be a_j
                // thus the sum is decreased by the difference between a_i and a_j 
                sum-=a[next]-(a[0]-d);
                next-=1;
            }
            else{
                // decrease the smallest element in the array by 1
                sum-=n-next; 
                d++;
            }
            steps++;
        }
        cout<<steps<<endl;
    }
}
 
/*
 
Hack test case:
 
10000
1 1
1000000000
1 1
1000000000
1 1
1000000000
1 1
1000000000
... (repeat this t=10000 times)
 
In this hack test case, my code will start at sum = 1 billion
and slowly decreases the sum by 1 for each iteration in the
while loop until the sum reaches 1.
 
estimated number of operations needed for my code to solve this test case:
a_i * t = 10^9 * 10*4 = 10^13 (takes > 1 hour)
 
*/

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int main(){
    ios::sync_with_stdio(0);cin.tie(0);
    ll t; cin>>t;
    for(ll i=0;i<t;i++){
        ll n; cin>>n;
        ll k; cin>>k;
        ll a[n];
        ll steps=0;
        ll sum=0;
        ll d=0;
        ll next=n-1;
        
        for(ll j=0;j<n;j++){
            cin>>a[j];
            sum+=a[j];
        }
        sort(a,a+n);
        while(sum>k){
            if(next==0){ // This change improves the code's time complexity.
                steps+=ceil((double)(sum-k)/n);
                sum=k-1;
            }
            
            else if(next>0 and a[next]-a[0]+d>n-next){
                sum-=a[next]-(a[0]-d);
                next-=1;
                steps++;
            }
            
            else{
                sum-=n-next; 
                d++;
                steps++;
            }
        }
        cout<<steps<<endl;
    }
}