Codeforces - C2. Potions (Hard Version)

 Problem Link

Algorithm used - Greedy, Data structures

Time Complexity: O(N log N)

You start out with a health value of 0. There are N potions numbered 1 to N that are examined sequentially. You may choose either drink or not drink a potion. Upon drinking a potion your health value changes by the potion's value. If the potion's is positive, then your health value increases by the potion's value. But if it's negative then your health value decreases by the potion's value. You must ensure that your health value remains non-negative. Find the maximum number of potions that you can drink.

It is obvious that you can drink all potions whose values are greater than or equal to 0, since your health value will never dip below 0 upon drinking these. To maximize the answer, we should also try to drink some negative potions.

To try and drink the maximum number of potions with a negative value, we sort them in ascending order and iterate over them. For each potion, check if our health value up until that potion is at least |negative potion|. (Use a segment tree for this since these will keep on changing) If it is, then we can drink it. After drinking this potion, our health value decreases. So, starting with the positive potion that is closest to the left of the current negative potion, keep on decrement the positive potion values in the segment tree until we have completely covered the cost of the negative potion. 

Code

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

const int MAX_N = 2 * 1e5;

ll nums[MAX_N];

vector <ll> seg_tree;

int seg_sz;

void update(int x, ll a) {

    x += seg_sz;

    seg_tree[x] += a;

    for (x /= 2; x >= 1; x /= 2)

        seg_tree[x] = seg_tree[x*2] + seg_tree[x*2+1];

}

ll sum(int a, int b) {

    a += seg_sz;

    b += seg_sz;

    ll s = 0;

    while (a <= b) {

        if (a % 2 == 1)

            s += seg_tree[a++];

        if (b % 2 == 0)

            s += seg_tree[b--];

        a /= 2;

        b /= 2;

    }

    return s;

}

int main() {

    int n;

    cin >> n;

    for (int i = 1; i <= n; i++)

        cin >> nums[i];

    seg_sz = pow(2, ceil(log2(n)));

    seg_tree.resize(seg_sz * 2);

    multiset <pair <ll, int>> negs;

    ll ans = 0;

    multiset <pair <ll, int>> pos;

    for (int i = 1; i <= n; i++) {

        if (nums[i] >= 0) {

            update(i-1, nums[i]);

            ans++;

            if (nums[i] > 0)

                pos.insert({i, nums[i]});

        } else

            negs.insert({-nums[i], i});

    }

    while (!negs.empty()) {

        ll x = negs.begin()->first;

        ll y = negs.begin()->second;

        if (sum(0, y-1) >= x) {

            ans++;

            while (x > 0) {

                auto a = (--pos.lower_bound({y, 0}))->first;

                if (nums[a] <= x) {

                    pos.erase(--pos.lower_bound({y, 0}));

                    update(a-1, -nums[a]);

                    x -= nums[a];

                    nums[a] = 0;

                } else {

                    pos.erase(--pos.lower_bound({y, 0}));

                    pos.insert({a, nums[a] - x});

                    nums[a] -= x;

                    update(a-1, -x);

                    x = 0;

                }

            }

        }

        negs.erase(negs.begin());

    }

    cout << ans << '\n';

}