Editorial for Cat-astrophe

Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Author: oree

In this problem, we need to calculate how many of the cat prices we can afford from left to right, for each of the thresholds. First of all, if we just loop through the prices until our sum exceeds the threshold for each threshold, we have $O(n\times q)$ ~O(n\times q)~. Both can be up to $10^5$ ~10^5~ so this would prove to be too slow, and we need a cleverer approach.

You can notice that a lot of the work done summing from left to right is repeated, and, when looking at any individual cat, we don't care about its price but moreso the price of it summed with all the cats to the left of it. Thus, we can construct a prefix sum over of the prices array with each element representing the price of that cat along with all the previous ones. Thus, we can see if we can afford up to any specific point in constant time, but looping through this array would still take $O(n\times q)$ ~O(n\times q)~. Instead, since this array is now non-decreasing, we can find the final cat we can afford using binary search. We can either manually implement binary search, or use the standard library std::upper_bound or bisect.bisect_right. Those would give you the index of the first greater element than the threshold, which, when considering 0-indexing, is the same as the number of cats we can afford.

In terms of time complexity, constructing the prefix sum takes $O(n)$ ~O(n)~, while each binary search takes $O(\log n)$ ~O(\log n)~, for a total time complexity of $O(n+q\log n)$ ~O(n+q\log n)~. The space complexity is $O(n)$ ~O(n)~ due to storing the prices and prefix sum.

C++ solution implementing binary search:

#include <iostream>
#include <vector>
using namespace std;

auto main() -> int {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int numCats, numQueries;
    cin >> numCats >> numQueries;

    vector<int> cats(numCats, 0);

    for (int i{0}; i < numCats; ++i) cin >> cats[i];

    vector<int> prefixSum(numCats, 0);
    prefixSum[0] = cats[0];
    for (int i{1}; i < numCats; ++i) {
        prefixSum[i] = cats[i] + prefixSum[i - 1];
    }

    while(numQueries--) {
        int thresholdQuery;
        cin >> thresholdQuery;

        auto withinBudget = [&prefixSum, &thresholdQuery](int mid) -> bool {
            if (prefixSum[mid] <= thresholdQuery) return true;
            else return false;
        };

        int left{-1};
        int right{numCats};

        while (right - left > 1) {
            int mid = left + (right - left) / 2;

            if (withinBudget(mid)) {
                left = mid;
            } else {
                right = mid;
            }
        }

        if (left <= -1 || left >= numCats) cout << 0 << '\n';
        else cout << left + 1 << '\n';
    }

    return 0;
}

C++ solution using std::upper_bound:

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n, q;
    cin >> n >> q;
    int curtot = 0;
    vector<int> prefix(n);
    for (auto& x: prefix) {
        cin >> x;
        x = (curtot += x);
    }

    int threshold;
    while (q--) {
        cin >> threshold;
        cout << upper_bound(prefix.begin(), prefix.end(), threshold)-prefix.begin() << endl;
    }
}

Python solution using bisect.bisect_right

from bisect import bisect_right

n, q = map(int, input().split())

prefix = list(map(int, input().split()))

for i in range(1, n):
    prefix[i] += prefix[i - 1]

for thresh in map(int, input().split()):
    print(bisect_right(prefix, thresh))

Comments

There are no comments at the moment.