It’s a very elegant problem. My solution executed in 0.04 seconds. The key to solving it is:
1 ≤ qi ( for i = 1 to Q ) ≤ 10000
This means NO query will be for something greater than 10000 position. Note, it is possible to generate 20000*20000 = 400000000 possible values. We are only interested in the first 10000 or so.
So here is how we proceed.
Step 1: Sort both Inputs.
Step 2: What is the minimum possible number that can be generated - A[0] +B[0]. Max is A[n-1]+B[n-1].
Step 3: Now for a given sum S, try to binary search down the possible ways to generate a number below it. So if we are trying to generate the number 100 or lesser, and the first five elements of array A are 10, 20, 30, 40, 50… we can find out the number of elements in array B less than 90,80,70,60 etc. Add these counts. The moment this count exceeds 10000, we can safely say that “100” is too large and lower our search range.
Step 4: Based on above analysis, let us say we have narrowed down 60 that produces over 10000 numbers - but not MUCH more than that. So now go ahead and generate those 10000 numbers and store them ina vector.
Step 5. Simply do offline queries to get answers.
Operative part of code below:
#define int long long int
pair<int,int> number_elements(int A[], int B[], int n, int qry, int targ, int old_targ)
{
int ans = 0;
int A_lt = lower_bound(A, A+n, targ) - A;
for (int i = 0; i < A_lt; ++i)
{
ans += lower_bound(B, B+n, targ-A[i]) - B;
if(ans > qry)
{
return (number_elements(A, B, n, qry, targ/2, targ));
break;
}
}
return(make_pair(targ, old_targ));
}
vector<int> enumerate_solve(int A[], int B[], int n, int targ, int qry)
{
vector<int> ans;
for(int i = 0; i < n; ++i)
{
if(A[i] > targ)
{
break;
}
for(int j = 0; j <n ; ++j)
{
if(A[i]+B[j] <= targ)
{
ans.push_back(A[i]+B[j]);
}
else
{
break;
}
}
}
sort(ans.begin(), ans.end());
return ans;
}
vector<int> solve(int A[], int B[], int n, int qry)
{
int minv = A[0] + B[0];
int maxv = A[n-1] + B[n-1];
int midv = minv + (maxv-minv)/2;
pair<int,int> ans = number_elements(A, B, n, qry, maxv, maxv);
return(enumerate_solve(A, B, n, ans.ss, qry));
}
int32_t main()
{
kpr; int t; cin >> t;
while(t--)
{
int k, q; cin >> k >> q; int A[k], B[k];
for (int i = 0; i < k; ++i)
{
cin >> A[i];
}
for (int i = 0; i < k; ++i)
{
cin >> B[i];
}
sort(A, A+k); sort(B, B+k);
int queries[q];
int maxQ = 0;
for (int i = 0; i < q; ++i)
{
cin >> queries[i];
maxQ = max(maxQ, queries[i]);
}
vector<int> precomp = solve(A, B, k, maxQ);
for (int i = 0; i < q; ++i)
{
print(precomp[queries[i]-1]);
}
}
return 0;
}