DQUERY - Editorial

PROBLEM LINK:

Contest
Practice

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;
        }
    }

}
2 Likes

Can someone explain how the pre comuputing part works?

let me explain you by an example
lets assume menu to be like this
[1,2,4,5,8,7]

prefix sum for above array : [1,3,7,12,20,27]
Let’s say p=2 and k=4
the optimal way of reordering would be [1,8,4,5,2,7]
for p=2 [2,4,8] are his favourite as 2 divides them
Here we must reorder items in such a way that deliousness is maximized
hence we’ll bring items which are maximum to front
the array sorted by values [8,4,2] pf sum : [8,12,14]

but we also have to remove sum of items present in the these places before reorder

hence we also need to sort them by indices
[2,4,8] pf sum arr: [2,6,14]

Now let’s calculate deliousness
firstly ,lets consider prefix sum of k items before reordering = 12

for accessing pf sum of fav items arrays we need no. of fav items present till [1,k]
which can be obtained for each query or technique discussed below

lets remove sum of those items which can be replaced with higher value = 2+4=6
lets add sum of those items which will be brought to front = 8+4=12

hence the answer can be calculated by 12-6+12==>18

technique to get no. of fav_items occured till k
we can store list of indices of fav_items and obtain lower bound of k in the list
hope this helps

3 Likes

Ahh…!!, Very Explanatory and Concise Editorial.

Thank you for the explanation, it cleared up many things.

1 Like