MXSZ - Editorial

editorial
, , , , , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Easy

PREREQUISITES:

Binary search, segment trees/fenwick trees

PROBLEM:

You’re given an array A of length N.
Answer Q queries.
Each query gives you X and Y, and you must report the largest size of a subsequence of elements from [A_1, A_2, \ldots, A_X] whose sum doesn’t exceed Y.

EXPLANATION:

First, we attempt to solve a single query.

We have the elements [A_1, \ldots, A_X], and we want to take as many of them as possible while not exceeding a sum of Y.
Clearly, this means it’s optimal to pick smaller elements first before larger ones.

More precisely, let S = \text{sorted}([A_1, \ldots, A_X]) denote the array obtained by sorting the first X elements of A.
Then, it’s optimal to choose some prefix of S, i.e. we’ll choose the elements S_1, S_2, \ldots, S_i for whichever is the largest index i satisfying S_1 + S_2 + \ldots + S_i \le Y.

Using the above knowledge, we now move to a slightly harder version of the problem - suppose the value of X is fixed across all queries, but Y can vary.

In this case, note that the array S remains constant.
We then only need to find, for a given Y, the largest index i such that S_1 + \ldots + S_i \le Y.
That is, we’re interested in the largest prefix sum of S that doesn’t exceed Y.

This can be easily done in \mathcal{O}(\log N) time with the help of binary search on the prefix sums.

So, assuming all queries have the same value of X, we’re able to solve the problem in \mathcal{O}(\log N) time per query after \mathcal{O}(N\log N) precomputation.

Finally, we move to the general version of the problem.

We will solve the queries offline, i.e. not in the order they’re given in the input.

First, let’s solve all queries that have X = N.
This is easy to do in \mathcal{O}(N\log N) as seen above: we only need to sort the entire array, then binary search on its prefix sums for each query.

After this, let’s try to solve all queries that have X = N-1.
Observe that not much actually changes when going from N to N-1.
In fact, the only change at all is that the element A_N gets removed from the sorted sequence - everything else remains the same!

So, if we are able to maintain a data structure that allows us to quickly remove an element from it, and also quickly tell prefix sums of existing elements, we would have enough power to solve the problem.

One way to do this is using a segment tree (with a little extra tracking).
Specifically, we will maintain the following information:

  • Let S = \text{sorted}(A).
  • Let P_i denote the position of element A_i in the sorted array S.
  • We also maintain an auxiliary array C, where C_i = 1 means the element at index i of S is still present.

We build a segment tree on the arrays S and C, storing range sums at each node.
Initially, C_i = 1 for all i.

Now, we process X = N, N-1, \ldots, 3, 2, 1 in order.
First, answer all queries with this value of X.

  • To answer a query with value Y, we can use binary search - specifically, binary search to find the largest index k such that S_1 + S_2 + \ldots + S_k \le Y.
  • The value of S_1 + \ldots + S_i for a fixed i can be found in \mathcal{O}(\log N) using the previously built segment tree.
  • Once the appropriate index k has been found, the answer is simply C_1 + C_2 + \ldots + C_k, denoting the number of ‘present’ elements among the considered prefix.

After all queries have been answered, we need to update our segment tree before reducing X further.
This is simple: the only change is that the value A_X is no longer available to us when moving to X-1.
So, we can do the following:

  • By definition, the value A_X appears at position P_X in the sorted array.
    So, we need to “disable” the contribution at index P_X.
  • This can be done by setting S_{P_X} = 0 and C_{P_X} = 0.

So, each query can be answered in \mathcal{O}(\log^2 N) by binary searching to find the appropriate prefix, and after each index we do a couple of point updates to keep the necessary data updates which takes \mathcal{O}(\log N) time.
This gives us a solution in \mathcal{O}(N\log N + Q\log^2 N) which is fast enough.

It’s possible to answer each query in \mathcal{O}(\log N) time by using the technique of segment tree walks, combining the binary search into the traversal of the tree to save a logarithmic factor.

TIME COMPLEXITY:

\mathcal{O}((N+Q)\log N) or \mathcal{O}(N\log N + Q\log^2 N) per testcase.

CODE:

Editorialist's code (C++) 
// #include <bits/allocator.h>
// #pragma GCC optimize("Ofast,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());

array<ll, 2> unit;
template<class T>
struct SegTree {
	T f(T a, T b) {
        return {a[0] + b[0], a[1] + b[1]};
    }
	vector<T> s; int n;
	SegTree(int _n = 0, T def = unit) : s(2*_n, def), n(_n) {}
	void update(int pos, T val) {
		for (s[pos += n] = val; pos /= 2;)
			s[pos] = f(s[pos * 2], s[pos * 2 + 1]);
	}
	T query(int b, int e) {
		T ra = unit, rb = unit;
		for (b += n, e += n; b < e; b /= 2, e /= 2) {
			if (b % 2) ra = f(ra, s[b++]);
			if (e % 2) rb = f(s[--e], rb);
		}
		return f(ra, rb);
	}
};


int main()
{
    ios::sync_with_stdio(false); cin.tie(0);

    int n, q; cin >> n >> q;
    vector<array<int, 3>> a(n);
    for (int i = 0; i < n; ++i) {
        cin >> a[i][0];
        a[i][1] = i;
    }

    ranges::sort(a);
    SegTree<array<ll, 2>> seg(n);
    for (int i = 0; i < n; ++i) {
        seg.update(i, {a[i][0], 1});
        a[i][2] = i;
    }
    ranges::sort(a, [&] (auto x, auto y) {return x[1] < y[1];});

    vector queries(n, vector<array<int, 2>>());
    for (int i = 0; i < q; ++i) {
        int x, y; cin >> x >> y;
        queries[--x].push_back({y, i});
    }

    vector ans(q, 0ll);
    for (int i = n-1; i >= 0; --i) {
        for (auto [y, id] : queries[i]) {
            int lo = 0, hi = n;
            while (lo < hi) {
                int mid = (lo + hi + 1)/2;
                if (seg.query(0, mid)[0] <= y) lo = mid;
                else hi = mid-1;
            }
            ans[id] = seg.query(0, lo)[1];
        }
        seg.update(a[i][2], unit);
    }

    for (auto x : ans) cout << x << ' ';
    cout << '\n';
}
1 Like