Power Products (Codeforces 1225D)
Tags: Math, hashing

You are given $$$n$$$ positive integers $$$a_1, \ldots, a_n$$$, and an integer $$$k \geq 2$$$. Count the number of pairs $$$i, j$$$ such that $$$1 \leq i < j \leq n$$$, and there exists an integer $$$x$$$ such that $$$a_i \cdot a_j = x^k$$$.

Input: The first line contains two integers $$$n$$$ and $$$k$$$ ($$$2 \leq n \leq 10^5$$$, $$$2 \leq k \leq 100$$$).The second line contains $$$n$$$ integers $$$a_1, \ldots, a_n$$$ ($$$1 \leq a_i \leq 10^5$$$).

Output: Print a single integer — the number of suitable pairs.

Note: In the sample case, the suitable pairs are: $$$a_1 \cdot a_4 = 8 = 2^3$$$; $$$a_1 \cdot a_6 = 1 = 1^3$$$; $$$a_2 \cdot a_3 = 27 = 3^3$$$; $$$a_3 \cdot a_5 = 216 = 6^3$$$; $$$a_4 \cdot a_6 = 8 = 2^3$$$.

Sample Inputs:
6 3
1 3 9 8 24 1

Sample Outputs:
5

Solution (Python):
import math as mt
from typing import Counter

def sieve():
    spf[1] = 1
    for i in range(2, MAXN):
        spf[i] = i
    for i in range(4, MAXN, 2):
        spf[i] = 2
    for i in range(3, mt.ceil(mt.sqrt(MAXN))):
        if spf[i] == i:
            for j in range(i * i, MAXN, i):
                if spf[j] == j:
                    spf[j] = i

def getFactorization(x, k):
    li = dict()
    while x != 1:
        if not spf[x] in li:
            li[spf[x]] = 0
        li[spf[x]] += 1
        if li[spf[x]] == k:
            del li[spf[x]]
        x = x // spf[x]
    li_prime = li.copy()
    for i in li_prime.keys():
        li_prime[i] = k - li_prime[i]
    return (tuple(li.items()), tuple(li_prime.items()))

def Solution(n, k, array, mask, mask_prime):
    lists = dict()
    sieve()
    count = 0
    for i in array:
        if len(mask[i]) == 0:
            (li, li_prime) = getFactorization(i, k)
            mask[i] = li
            mask_li_prime[i] = li_prime
        if mask_li_prime[i] in lists:
            count += lists[mask_li_prime[i]]
        if mask_li[i] in lists:
            lists[mask_li[i]] += 1
        else:
            lists[mask_li[i]] = 1
    return count
values = input().split()
(n, k) = (int(values[0]), int(values[1]))
a = list(map(int, input().split()))
maximun = max(a)
MAXN = maximun + 1
mask_li = [() for i in range(0, maximun + 1)]
mask_li_prime = [() for i in range(0, maximun + 1)]
spf = [0 for i in range(MAXN)]
result = Solution(n, k, a, mask_li, mask_li_prime)
print(result)

Solution (Python):
import math as mt
from typing import Counter

def sieve():
    spf[1] = 1
    for i in range(2, MAXN):
        spf[i] = i
    for i in range(4, MAXN, 2):
        spf[i] = 2
    for i in range(3, mt.ceil(mt.sqrt(MAXN))):
        if spf[i] == i:
            for j in range(i * i, MAXN, i):
                if spf[j] == j:
                    spf[j] = i

def getFactorization(x, k):
    li = dict()
    while x != 1:
        if not spf[x] in li:
            li[spf[x]] = 0
        li[spf[x]] += 1
        if li[spf[x]] == k:
            del li[spf[x]]
        x = x // spf[x]
    li_prime = li.copy()
    for i in li_prime.keys():
        li_prime[i] = k - li_prime[i]
    return (tuple(li.items()), tuple(li_prime.items()))

def Solution(n, k, array):
    lists = dict()
    sieve()
    count = 0
    for i in array:
        if len(mask_li[i]) == 0:
            (li, li_prime) = getFactorization(i, k)
            mask_li[i] = li
            mask_li_prime[i] = li_prime
        if mask_li_prime[i] in lists:
            count += lists[mask_li_prime[i]]
        if mask_li[i] in lists:
            lists[mask_li[i]] += 1
        else:
            lists[mask_li[i]] = 1
    return count
values = input().split()
(n, k) = (int(values[0]), int(values[1]))
a = list(map(int, input().split()))
maximun = max(a)
MAXN = maximun + 1
mask_li = [() for i in range(0, maximun + 1)]
mask_li_prime = [() for i in range(0, maximun + 1)]
spf = [0 for i in range(MAXN)]
result = Solution(n, k, a)
print(result)

Solution (Python):
(n, k, *a) = map(int, open(0).read().split())
p = [*range(7 ** 6)]
for i in range(2, 7 ** 6):
    if p[i] == i:
        for j in range(i, 7 ** 6, i):
            p[j] = i
ans = 0
mp = {}
for x in a:
    F = {}
    while x > 1:
        f = p[x]
        x //= f
        F[f] = F.get(f, 0) + 1
        if F[f] == k:
            del F[f]
    F = [*F.items()]
    T = tuple(F)
    for (i, (f, c)) in enumerate(F):
        c = k - c
        if c == k:
            c = 0
        F[i] = (f, c)
    ans += mp.get(tuple(F), 0)
    mp[T] = mp.get(T, 0) + 1
print(ans)