PROBLEM LINK:
Setter: rag_hav13
Testers: tabr
Editorialist: aryanag_adm
DIFFICULTY:
2484
PREREQUISITES:
Factorisation, Prefix Sums, Binary Search
PROBLEM:
You are given a list of N integers A_1, A_2, \cdots, A_n, along with Q queries.
Each query consists of two integers p and k.
Suppose you are allowed to reorder only those elements of A which are divisible by p.
You have to output the maximum possible value of \sum_{i=1}^k A_i over all such reorderings.
EXPLANATION:
Our reordering strategy is trivial.
Let S = \{A_i | \ i \in [N], A_i \text{ is divisible by } p\}. We would like to place the largest element from S at the smallest index, the second largest element at the second smallest index, so on and so forth. \
What would be the value of \sum_{i=1}^k A_i in such a reordering?
Define T = \{i \in [k] \ | \ A_i \text{ is divisible by } p\}.
Then, in our reordering, we would place the largest |T| elements from S into indices in T, and we would place the elements that are already in those T indices somewhere else (possibly outside the range [1, k]).
Therefore, the answer we get would be \sum_{i=1}^k A_i + (\text{sum of the largest } |T| \text{ elements of }S) - \sum_{i \in T} A_i
We can precompute the factorisation of all numbers in range [1, 1e5], and use this to prime factorise M. Note that each element of M can have at most \log 1e5 < 17 prime factors. Now, for each prime number \in [1, 1e5], we can precompute S. We can precompute suffix sums of S (when sorted by value) and prefix sums of S (when sorted by index). Finally, we can precompute prefix sums of M.
Now we describe how to answer each query. We can find |T| by doing a simply binary search for k in S, this takes O(\log n). We can compute \sum_{i=1}^k A_i in O(1) using the prefix sums we precomputed. We can compute (\text{sum of the largest } |T| \text{ elements of }S) in O(1) using the suffix sums we precomputed. And finally, we can compute \sum_{i \in T} A_i using the prefix sums of S in O(1), that we also precomputed.
TIME COMPLEXITY:
Approximately O(M\log M + 1e5\log1e5) precomputation, and then O(\log N) per query, where M = N \log 1e5.
SOLUTION:
Editorialist's Solution
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
//#include <sys/resource.h>
#define double long double
#define int long long
#define initrand mt19937 mt_rand(time(0));
#define rand mt_rand()
#define MOD 1000000007
#define INF 1000000000
#define mid(l, u) ((l+u)/2)
#define rchild(i) (i*2 + 2)
#define lchild(i) (i*2 + 1)
#define mp(a, b) make_pair(a, b)
#define lz lazup(l, u, i);
#define ordered_set tree<pair<int, int>, null_type,less<pair<int, int>>, rb_tree_tag,tree_order_statistics_node_update>
using namespace std;
using namespace __gnu_pbds;
bool notPrime[100001];
vector<int> pfac[100001];
signed main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
cout.tie(NULL);
for(int i = 2;i<=100000;i++){
if(notPrime[i]) continue;
for(int j = i;j<=1e5;j+=i){
notPrime[j] = true;
pfac[j].push_back(i);
}
}
int t;
cin>>t;
while(t--) {
int n;
cin >> n;
int m[n];
for(int i = 0;i<n;i++) cin>>m[i];
int pref[n];
pref[0] = m[0];
for(int i = 1;i<n;i++) pref[i] = m[i] + pref[i-1];
vector<int> li[100001];
for(int i = 0;i<n;i++){
for(int j: pfac[m[i]]){
li[j].push_back(i);
}
}
vector<int> pf[100001], sf[100001];
for(int i = 0;i<=100000;i++){
if(li[i].size() == 0) continue;
pf[i].push_back(m[li[i][0]]);
for(int j = 1;j<li[i].size();j++){
pf[i].push_back(m[li[i][j]] + pf[i][j-1]);
}
vector<int> temp;
for(int j: li[i]) temp.push_back(m[j]);
sort(temp.begin(), temp.end());
sf[i].push_back(temp[temp.size() - 1]);
for(int j = 1;j<temp.size();j++){
sf[i].push_back(temp[temp.size() - 1 - j] + sf[i][j-1]);
}
}
int q;
cin>>q;
while(q--){
int p, k;
cin>>p>>k;
int ans = pref[k-1];
int cnt = lower_bound(li[p].begin(), li[p].end(), k) - li[p].begin();
if(cnt > 0){
ans -= pf[p][cnt - 1];
ans += sf[p][cnt - 1];
}
cout<<ans<<endl;
}
}
}