Official Solution

This uses merge-sort tree.

Time Complexity: $O(N\log N+Q\log^2N)$

Memory Complexity: $O(N\log N)$

Alternative solution (Wavelet Tree)

Similar as K-query from SPOJ, but you additionally need to store prefix sums of $A$.

Time Complexity: $O((N+Q)\log M)$, where $M$ is the max value

Memory Complexity: $O(N\log M)$

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

struct PrefixSummer {
	const int BITS = 64;
	vector<uint64_t> packed;
	vector<int> psums;
	vector<int64_t> psumA;
	void init(const vector<bool> &v, vector<int> A) {
		packed.resize(size(v) / BITS + 1);