Count pairs in Array whose product is a Kth power of any positive integer

Given an array arr[] of length N and an integer K, the task is to count pairs in the array whose product is Kth power of a positive integer, i.e.

A[i] * A[j] = Z^K for any positive integer Z.

Examples:

Input: arr[] = {1, 3, 9, 8, 24, 1}, K = 3 

Output: 5 

Explanation: There are 5 such pairs, those can be represented as Z³ – A[0] * A[3] = 1 * 8 = 2^3 A[0] * A[5] = 1 * 1 = 1^3 A[1] * A[2] = 3 * 9 = 3^3 A[2] * A[4] = 9 * 24 = 6^3 A[3] * A[5] = 8 * 1 = 2^3 

Input: arr[] = {7, 4, 10, 9, 2, 8, 8, 7, 3, 7}, K = 2 

Output: 7 

Explanation: There are 7 such pairs, those can be represented as Z²

Approach: The key observation in this problem is for representing any number in the form of Z^K then powers of prime factorization of that number must be multiple of K. Below is the illustration of the steps:

• Compute the prime factorization of each number of the array and store the prime factors in the form of key-value pair in a hash-map, where the key will be a prime factor of that element and value will be the power raised to that prime factor modulus K, in the prime factorization of that number. 

For Example:

Given Element be - 360 and K = 2
Prime Factorization = 23 * 32 * 51

Key-value pairs for this would be,
=> {(2, 3 % 2), (3, 2 % 2),
    (5, 1 % 2)}
=> {(2, 1), (5, 1)}

// Notice that prime number 3 
// is ignored because of the 
// modulus value was 0

• Traverse over the array and create a frequency hash-map in which the key-value pairs would be defined as follows:

Key: Prime Factors pairs mod K
Value: Frequency of this Key

• Finally, Traverse for each element of the array and check required prime factors are present in hash-map or not. If yes, then there will be F number of possible pairs, where F is the frequency. 

Example:

Given Number be - 360, K = 3
Prime Factorization -
=> {(3, 2), (5, 1)}

Required Prime Factors -
=> {(p1, K - val1), ...(pn, K - valn)}
=> {(3, 3 - 2), (5, 3 - 1)}
=> {(3, 1), (5, 2)}  

Below is the implementation of the above approach: 

C++

// C++ implementation to count the
// pairs whose product is Kth
// power of some integer Z

#include <bits/stdc++.h>

#define MAXN 100005

using namespace std;

// Smallest prime factor
int spf[MAXN];

// Sieve of eratosthenes
// for computing primes
void sieve()
{
    int i, j;
    spf[1] = 1;
    for (i = 2; i < MAXN; i++)
        spf[i] = i;

    // Loop for markig the factors
    // of prime number as non-prime
    for (i = 2; i < MAXN; i++) {
        if (spf[i] == i) {
            for (j = i * 2;
                 j < MAXN; j += i) {
                if (spf[j] == j)
                    spf[j] = i;
            }
        }
    }
}

// Function to factorize the
// number N into its prime factors
vector<pair<int, int> > getFact(int x)
{
    // Prime factors along with powers
    vector<pair<int, int> > factors;

    // Loop while the X is not
    // equal to 1
    while (x != 1) {

        // Smallest prime
        // factor of x
        int z = spf[x];
        int cnt = 0;
        // Count power of this
        // prime factor in x
        while (x % z == 0)
            cnt++, x /= z;

        factors.push_back(
            make_pair(z, cnt));
    }
    return factors;
}

// Function to count the pairs
int pairsWithKth(int a[], int n, int k)
{

    // Precomputation
    // for factorisation
    sieve();

    int answer = 0;

    // Data structure for storing
    // list L for each element along
    // with frequency of occurrence
    map<vector<pair<int,
                    int> >,
        int>
        count_of_L;

    // Loop to iterate over the
    // elements of the array
    for (int i = 0; i < n; i++) {

        // Factorise each element
        vector<pair<int, int> >
            factors = getFact(a[i]);
        sort(factors.begin(),
             factors.end());

        vector<pair<int, int> > L;

        // Loop to iterate over the
        // factors of the element
        for (auto it : factors) {
            if (it.second % k == 0)
                continue;
            L.push_back(
                make_pair(
                    it.first,
                    it.second % k));
        }

        vector<pair<int, int> > Lx;

        // Loop to find the required prime
        // factors for each element of array
        for (auto it : L) {

            // Represents how much remainder
            // power needs to be added to
            // this primes power so as to make
            // it a multiple of k
            Lx.push_back(
                make_pair(
                    it.first,
                    (k - it.second + k) % k));
        }

        // Add occurrences of
        // Lx till now to answer
        answer += count_of_L[Lx];

        // Increment the counter for L
        count_of_L[L]++;
    }

    return answer;
}

// Driver Code
int main()
{
    int n = 6;
    int a[n] = { 1, 3, 9, 8, 24, 1 };
    int k = 3;

    cout << pairsWithKth(a, n, k);
    return 0;
}

Output:

5

Time complexity: O(N * log²N)