AATI2407 - Editorial

editorial
,
1

Problem Link:

Contest
Practice

Author: Pankaj Goyal
Tester: Chaitanya Hardikar
Editorialist: Pankaj Goyal

Difficulty

Medium

PREREQUISITES

Sieve of Eratosthenes
Observation

EXPLANATION & SOLUTION:

Firstly, we will pre calculate the sum of all the factors so that we can directly use the results for all of the test cases. Using these results, we can construct the array B.
For the second part we shall simplify the conditions a bit more so that we can process it in a better way. We need to find i, j such that i < j and B[i]+j-i = B[j]+i-j

Let’s analyze the second condition.
B[i]+j-i = B[j] + i - j
B[i] - i - i = B[j] - j - j
B[i] - 2*i = B[j] - 2*j

Let’s call B[i]-2*i as value of index i. Therefore, our problem reduces to finding the number of indices whose value are equal. This can be efficiently done using a map.
Let’s say if there are ‘x’ indices of same value, then its contribution as a pairwise will be x*(x-1)/2. This way M1 is calculated.

For calculating M2 and M3 we just need to find which indices are contributing to the total pairwise sum. Now since we mapped all the values of indices to the map, we will traverse the entire array and check if that particular index’s value is used or not by checking the mapped value if it’s greater than 1. This way array C is constructed. Now the count of primes in the array C
would be M2 and non-primes would be M3. Checking primes can be efficiently done using the Sieve of Eratosthenes.

Note: Calculating the sum of factors can be done in two ways, either by prime factorization or by using the multiples. In case of prime factorization, the time to find so would be O(√x) where x is the given number and therefore total time complexity would be O(N√N) where N is 10^6. Other way be just going through all the multiples and adding the contribution to it. Time complexity for this would be as follows

O (N/1 + N/2 + N/3 + …. + N/N)
O (N \{1 + ½ + 1/3 + …+ 1/N\})
O (N log(N))

Time complexity: O (x log(x) + N) where x is 10^6.

Setter's Solution 
#include <bits/stdc++.h>
using namespace std;
#define ll long long
int main() {
// Generating the sum of factors which will be common for all the test cases
vector < int > sum_of_factors(1e6 + 1, 0);
for (int i = 1; i <= 1e6; i++) {
  for (int j = i; j <= 1e6; j += i)
    sum_of_factors[j] += i;
}
vector < bool > prime(1e6 + 1, true);
prime[0] = prime[1] = false;
// Using the seive to find all the primes below 10^6
for (int i = 2; i * i <= 1e6; i++) {
  if (prime[i]) {
    for (int j = i * i; j <= 1e6; j += i)
      prime[j] = false;
  }
}
int t;
cin >> t;
while (t--) {
  int n;
  cin >> n;
  vector < int > A(n);
  for (int i = 0; i < n; i++)
    cin >> A[i];
  // Generating the array B
  vector < int > B(n);
  for (int i = 0; i < n; i++)
    B[i] = sum_of_factors[A[i]];
  // we have B[i]+j-i=B[j]+i-j which simplifies to
  // B[i] - 2 * i = B[j] - 2 * j
  /* Therefore we will map the values of B[ind] -
  2 * ind and the choosing all possible pairs as value choose
  2 will
  return all the pairs satisfying the condition
  Mapping the values to the map which we will
  be using to find value of M1  */
  map < int, ll > mp;
  for (int i = 0; i < n; i++) mp[B[i] - 2 * i]++;
  ll M1 = 0, M2 = 0, M3 = 0;
  /* map all the values of freq*(freq-1)/2 which are
  greater than so that we can separate all corresponding
  B[i]'s */
  map < int, ll > check;
  for (auto it: mp) {
    M1 += (it.second * (it.second - 1)) / 2;
    if (it.second > 1)
      check[it.first] = 1;
  }
  /* Traverse through the array B and check if its
  contributing or not and accordingly construct array C */
  vector < int > C;
  for (int i = 0; i < n; i++) {
    if (check.find(B[i] - 2 * i) != check.end())
      C.push_back(B[i]);
  }
  // Traverse the array C and calculate M2 and M3 using the precalculated primes array
  for (int i = 0; i < (int) C.size(); i++) {
    if (prime[C[i]])
      M2 += C[i];
    else
      M3 += C[i];
  }
  cout << M1 << " " << M2 << " " << M3 << "\n";
}
return 0;
}
1 Like