倍数遍历: 代码源每日一题-平方计数

通过倍数遍历降低时间复杂度， 题目链接

题目

题目描述

给出一个长度为$n$的正整数数组$a$， 问有多少对$(i, j)$满足$a_i^2 + a_j$是一个完全平方数

数据范围

$1 \leq n, a_i \leq 10^6$

题解

思路

$a_i^2 + a_j$是一个平方数等价于

$$a_i^2 + a_j = x^2$$

移项

$$a_j = x^2 - a_i^2$$

变形

$$a_j = (x - a_i) \times (x + a_i)$$

显然$x - a_i, x + a_i$都是$a_j$的约数， 那我们不妨枚举$c = x - a_i$为$a_j$的约数， 令$d = x + a_i = a_j / c$， 则$a_i = (d - c) / 2$， 然后就可以统计出答案了

但是对于每个$a_j$求一遍约数是$O(\sqrt n)$的， $n$个数就是$O(n \sqrt n)$， 过不了

倍数遍历优化

对多个数每个数都求约数， 数的数量多但是都比较小， 这个时候就可以枚举约数然后用约数的倍数遍历每个数， 时间复杂度为

$$\frac{maxnum}{1} + \frac{maxnum}{2} + \frac{maxnum}{3} + \dots + \frac{maxnum}{maxnum} = O(n \ln n)$$

比常规的$O(n \sqrt n)$快很多

代码

#include <iostream>
#include <cstdio>

int n;
int cnt[1123456];
long long ans;

int main()
{
    scanf("%d", &n);
    for(int i = 1; i <= n; i ++)
    {
        int a;
        scanf("%d", &a);

        cnt[a] ++;
    }

    for(int i = 1; i <= 1000000; i ++) // i为约数
        for(int j = i; j <= 1000000; j += i) // j为aj
            if((j / i - i) > 0 && (j / i - i) % 2 == 0)
                ans += 1LL * cnt[j] * cnt[(j / i - i) / 2];

    printf("%lld\n", ans);

    return 0;
}