[Acwing4319] 合适数对

数论算法同余

给定一个长度为 $n$ 的正整数数列 $a_1,a_2,…,a_n$ 和一个正整数 $k$ 。

请你判断共有多少个数对 $(l,r)$ 同时满足：

$1≤l<r≤n$

存在一个整数 $x$ 使得 $a_l×a_r=x^k$ 成立

输入格式

第一行包含两个整数 $n,k$ 。

第二行包含 $n$ 个正整数 $a_1,a_2,…,a_n$ 。

输出格式

一个整数，表示满足条件的数对的数量。

数据范围

前三个测试点满足 $2≤n≤10$ 。
所有测试点满足 $2≤n≤10^5,2≤k≤100,1≤a_i≤10^5$ 。

Acwing上的一道Hard难度的周赛题。

首先我们考察这个条件：存在一个整数 $x$ 使得 $a_l×a_r=x^k$ 成立。注意到 $a_l,a_r>0$ ，因此不妨认为 $x$ 是正整数。

由算术基本定理，任意一个正整数都能唯一地分解为全体质数的幂积：

N = \prod {p_{i}}^{m_{i}}, 其 中 N > 1, p_{i} 是 从 小 到 大 排 列 的 互 不 相 同 的 全 部 质 数, m_{i} 是 非 负 整 数

$N = \prod {p_i}^{m_i},其中 N>1,p_i是从小到大排列的互不相同的全部质数,m_i是非负整数$

于是，我们可以根据上式将任意一个正整数表示为一个无限维行向量：

\vec{N} = (m_{1}, m_{2}, m_{3}, \dots)

$\vec{N} = (m_1,m_2,m_3,\dots)$

例如， $42 = 2^13^17^1 = 2^13^15^07^19^0\dots$ ，那么， $\vec{42} = (1,1,0,1,0,0,0,\dots)$ 。特别地 $\vec{1}= (0,0,0,\dots)$

那么，任意两个正整数的乘法运算就可以转换为向量加法运算：

a \times b = \vec{a} + \vec{b}

$a\times b = \vec{a}+\vec{b}$

于是我们知道 $\vec{x^k} = k\vec{x}$ ，向量 $\vec{x^k}$ 的每一个元素都能被 $k$ 整除

那么，我们需要找出的满足 $a_l\times a_r = x^k$ 的正整数 $a_l,a_r$ 满足性质： $\vec{a_l} 和 \vec{a_r}中的对应元素之和都能被k整除$

注意到我们并不关心 $\vec{a_l}+\vec{a_r}$ 的每个元素的具体值，只关心它们是否都能被 $k$ 整除。那么，接下来就应该考虑什么情况下两数之和能被 $k$ 整除。

由同余代数的知识可以知道： $k$ 的任意整数倍模 $k$ 都是 $0$ ，如果 $a \equiv c \bmod k$ ， $b \equiv d \bmod k$ ，那么 $a+b \equiv c+d \bmod k$ ，那么，如果要有 $a+b$ 能被 $k$ 整除，一定有 $c+d \equiv 0 \bmod k$

设 $-d \equiv d_补 \bmod k$ ，那么 $c \equiv d_补 \bmod k$ ，即 $c \equiv -d \bmod k$

于是，如果我们知道一个数 $d$ ，那么所有与 $d$ 加起来能被 $k$ 整除的数都与 $-d \mod k$ 同余

于是乎，所有与 $\vec{a}$ 加起来后每一个元素能被 $k$ 整除的向量 $\vec{b}$ ，一定满足：

b_{i} = - a_{i} \mod k, b_{i} 与 a_{i} 是 向 量 \vec{b} 和 \vec{a} 的 对 应 位 置 上 的 元 素

$b_i = -a_i \mod k, b_i与a_i是向量\vec{b}和 \vec{a}的对应位置上的元素$

于是，我们就有了解题思路：

从前往后枚举每一个元素 $a_i$ ，利用分解质因数算法计算 $a_i$ 的向量，求出该向量每个元素模 $k$ 后的向量 $A_{ik}$ 和每个元素取负后再模 $k$ 的向量 $(-A_i)_k$
在一个记录向量出现次数的数据结构中寻找向量 $(-A_i)_k$ 的出现次数，将它累加到答案
将 $(A_i)_k$ 记录到数据结构，即其出现次数加1

注意到计算向量会有些复杂，因为我们需要知道每个质数在向量中的位置，由于一个正整数和它的向量是唯一对应的，于是在分解质因数时，我们只计算向量 $(A_i)_k$ 与 $(-A_i)_k$ 对应的正整数值 $(a_i)_k$ 和 $(-a_i)_k$ 即可。

由于每个质数因子的指数取模后一定不会变大，一定有 $(a_i)_k \leq a_i$ 。而取负后再取模就很可能变得更大，比如 $-1 \mod 100 = 99$ ，因此 $(-a_i)_k$ 有可能大于 $10^5$ ，这种情况下一定没有 $(a_j)_k = (-a_i)_k$ ，要注意判断。

分解质因数的时间复杂度为 $O(\sqrt n)$ ~ $O(\log n)$ ，总时间复杂度为 $O(n\sqrt n)$ ~ $O(n \log n)$

#include<iostream>
#include<unordered_map>
#include<string>
#include<vector>
#include<cmath>
using namespace std;
typedef unsigned long long ULL;
typedef pair<ULL,ULL> PUU;

const int N = 100010;

PUU K_num_anti(int x, int k){
    PUU ans={1,1};
    for(int i = 2;i<=x/i;i++){
        if(x%i == 0){
            int t = 0;
            while(x%i == 0){
                t++;
                x/=i;
            }
            //求幂
            for(int j = 1;j<= (t%k);j++){
                ans.first *= i;
            }
            for(int j = 1;j<=((-t)%k+k)%k;j++){
                if(ans.second>N){ans.second = N;break;}
                ans.second*= i;
            }
        }
    }
    if(x>1){
        ans.first *= x;
        for(int i = 1;i<=((-1)%k+k)%k;i++){
            if(ans.second>N){ans.second = N;break;}
            ans.second*= x;
        }
    }
    return ans;
}

unordered_map<ULL,int> fd;
int p[N];
int main(){
    int n,k;
    scanf("%d%d", &n,&k);
    for(int i = 1;i<=n;i++)scanf("%d", p+i);
    
    ULL ans = 0;

    for(int i = 1;i<=n;i++){
        int num = p[i];
        // res 第一个存[ai]k，第二个存 [-ai]k,如果它是大于N的，存N
        auto res = K_num_anti(num,k);
        
        if(res.second!=N && fd.find(res.second)!=fd.end())ans += fd[res.second];
        if(fd.find(res.first) != fd.end()) fd[res.first]++;
        else fd.emplace(res.first,1);
    }
    printf("%lld\n",ans);
    return 0;
}