Math, Implementation

First check if the sum of the array is divisible by its size.

If not, think about removing one element and what happens to the sum and the size.

The element you remove must have the right remainder when divided by (n−1).

We are given an array of size n with total sum S. If S is divisible by n, the array is already balanced and the answer is "YES". Otherwise, we can try deleting one element ai. After deletion, the new sum becomes S — ai and the new size becomes n — 1. For the new array to be balanced, (S — ai) must be divisible by (n — 1), which is the same as saying that ai must have the same remainder as S when divided by (n — 1). Thus, it is enough to check if there exists at least one element such that ai % (n — 1) == S % (n — 1). A special case is when n = 1, where the answer is always "YES" since deleting the single element leaves an empty array, which is considered balanced. The solution only requires scanning the array once, giving a time complexity of O(n) per test case.

~~~~~

~~~~~

Sorting + greedy

We always want to choose the cheapest item available.

Think of data structure used to maintain the cheapest item.

Use Minheap

Push all the elements into a minheap while also storing their index. Then pop out the cheapest fruit and we buy it (decrease the remaining budget). Then we push the updated price of the fruit back into the heap. The updated price = price of fruit + index.

We repeat the above process till the cheapest fruit that we get is costlier than the remaining budget.

Time Complexity :- Initially one might think that we can get TLE because budget <= 1e9.

But lets consider a worst case scenario :- B = 1e9 N = 1 A[1] = 1 In this case we add in the series 1 + 2+ 3 + 4 + — — — — +x such that x is the last price at which we buy.

Here we can use the formula x*(x+1)/2 <= B.

So approximately x = sqrt(B) ~ sqrt(1e9) ~ 1e5

Also x is the number of times we pop an element from the heap and add another so the final time complexity :-

O(sqrt(B)*log(n)).

Space complexity = O(n). We only need to use the space used in the heap.

#include <iostream>
#include <vector>
#include <queue>

using namespace std;

int main() {
    /* I might use long long when int works and I suggest at least when learning competitive programming you should ignore space complexity most of the time.
                                            Not that it is not important but currently let's focus on problem solving. */

    long long B;
    cin >> B;

    int n;
    cin >> n;

    vector<long long> a(n); // using long long for safety
    for (int i = 0; i < n; ++i) {
        cin >> a[i];
    }

    // Min heap declaration
    priority_queue<pair<long long, int>, vector<pair<long long, int>>, greater<pair<long long, int>>> pq;

    for (int i = 0; i < n; ++i) {
        pq.push({a[i], i + 1}); // pushing all the elements
    }

    long long ans = 0;
    while (!pq.empty()) {
        if (B < pq.top().first) { // base condition when cheapest > remaining budget
            break;
        } else {
            long long cur = pq.top().first; // get the cheapest fruit
            long long ind = pq.top().second; // get the index
            pq.pop(); // remove the chosen fruit

            B -= cur; // decrease budget
            ++ans; // One more fruit bought

            pq.push({cur + ind, ind}); // Update the minheap with new value
        }
    }

    cout << ans << endl; // Done output the answer

    return 0;
}

~~~~~

~~~~~

~~~~~

~~~~~

~~~~~

~~~~~