[Acwing4319] 合适数对

[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是非负整数 \]

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

\[\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} \]

于是我们知道 \(\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}的对应位置上的元素 \]

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

1. 从前往后枚举每一个元素 \(a_i\)，利用分解质因数算法计算 \(a_i\) 的向量，求出该向量每个元素模 \(k\) 后的向量 \(A_{ik}\) 和每个元素取负后再模 \(k\) 的向量 \((-A_i)_k\)

2. 在一个记录向量出现次数的数据结构中寻找向量 \((-A_i)_k\) 的出现次数，将它累加到答案

3. 将\((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;
}