1622C - Set or Decrease

A few others and I had solutions for this that passed the system tests with a bad time complexity of O(a_i * t). The problem has a_i <= 10^9 and t <= 10^4 so a_i * t can be up to 10^13, which could take over an hour for the code to run.

For any step, if you choose "set", it is optimal to set the greatest element (that hasn't been set yet) equal to the smallest element, because this decreases the sum the most.

If you choose "decrease", it is optimal to decrease the smallest element.

It's optimal to do all the "decrease"s before all the "set"s, because this way the largest elements can be set to an even lower number, compared to if any "set" step is done before a "decrease" step.

My approach (with the wrong time complexity) was to simulate the steps one by one, on each step choosing greedily whether to set or decrease based on which would decrease the sum the most.

(C++20)

#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)
 
*/

Generator for a test case that hacks the above code with "Time Limit Exceeded" verdict: (Python 3)

print(10000)
for i in range(10000):
    print("1 1")
    print("1000000000")

If all the other elements are already going to be "set" to the smallest element, the optimal move is always to choose "decrease". With this insight, we can break out of the while loop early in this situation, which happens in the 1 1 1000000000 test case.

Fixed code with correct time complexity: 140835404