Power Products (Codeforces 1246B)
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):
from collections import Counter as C
NN = 100000
X = [-1] * (NN + 1)
L = [[] for _ in range(NN + 1)]
k = 2
while k <= NN:
    X[k] = 1
    L[k].append(k)
    for i in range(k * 2, NN + 1, k):
        X[i] = 0
        L[i].append(k)
    d = 2
    while k ** d <= NN:
        for i in range(k ** d, NN + 1, k ** d):
            L[i].append(k)
        d += 1
    while k <= NN and X[k] >= 0:
        k += 1
P = [i for i in range(NN + 1) if X[i] == 1]
K = 4

def free(ct):
    a = 1
    for c in ct:
        a *= c ** (ct[c] % K)
    return a

def inv(ct):
    a = 1
    for c in ct:
        a *= c ** (-ct[c] % K)
    return a
(N, K) = map(int, input().split())
A = [int(a) for a in input().split()]
D = {}
ans = 0
for a in A:
    c = C(L[a])
    invc = inv(c)
    if invc in D:
        ans += D[invc]
    freec = free(c)
    if freec in D:
        D[freec] += 1
    else:
        D[freec] = 1
print(ans)

Solution (Python):
from collections import defaultdict
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def fast_prime_factorization(n):
    d = [(i + 1) % 2 * 2 for i in range(n + 1)]
    (d[0], d[1]) = (0, 1)
    for i in range(3, n + 1):
        if d[i]:
            continue
        for j in range(i, n + 1, 2 * i):
            if not d[j]:
                d[j] = i
    return d
(n, k) = map(int, input().split())
a = list(map(int, input().split()))
l = pow(10, 5) + 5
d = fast_prime_factorization(l)
cnt = defaultdict(lambda : 0)
for i in a:
    j = i
    s = defaultdict(lambda : 0)
    while j ^ d[j]:
        s[d[j]] += 1
        j //= d[j]
    if j ^ 1:
        s[j] += 1
    u = list(s.keys())
    u.sort()
    lu = len(u)
    v = []
    for i in u:
        if s[i] % k:
            v.append(i)
            v.append(s[i] % k)
    cnt[tuple(v)] += 1
ans0 = 0
ans = 0
for u in cnt.keys():
    v = list(u)
    w = list(v)
    f = 0
    for i in range(1, len(v), 2):
        v[i] = -v[i] % k
        if v[i] ^ w[i]:
            f = 1
    v = tuple(v)
    if f:
        if v in cnt:
            ans0 += cnt[u] * cnt[v]
    else:
        c = cnt[u]
        ans += c * (c - 1) // 2
ans += ans0 // 2
print(ans)

Solution (Python):
nmax = 400
eratos = [0 for i in range(nmax + 1)]
prime = []
cnt = 2
while True:
    while cnt <= nmax and eratos[cnt]:
        cnt += 1
    if cnt > nmax:
        break
    eratos[cnt] = 1
    prime.append(cnt)
    for i in range(cnt ** 2, nmax + 1, cnt):
        eratos[i] = 1
from collections import defaultdict
import sys
input = sys.stdin.readline
(n, k) = map(int, input().split())
a = list(map(int, input().split()))
st = set()
for i in prime:
    if i ** k > 10 ** 5:
        break
    st.add(i ** k)
dc = defaultdict(int)
ans = 0
for i in a:
    for j in st:
        while i % j == 0:
            i //= j
    dc[i] += 1
for (i, v) in dc.items():
    x = 1
    y = i
    if i == 1:
        ans += (v - 1) * v // 2
        continue
    if v == 0:
        continue
    for p in prime:
        cnt = 0
        while i % p == 0:
            cnt += 1
            i //= p
        if cnt:
            x *= p ** (k - cnt)
    if i != 1:
        if k == 2:
            x *= i
        else:
            continue
    if x == y:
        ans += v * (v - 1) // 2
        continue
    if x in dc:
        ans += dc[x] * v
        dc[x] = 0
print(ans)