2018ICPC北京 C. Pythagorean triple

Oct 23, 2019 由 小羊

数出所有满足 $a^2+b^2=c^2$ 且 $c \le n$ 的正整数三元组。

首先，如果 $(x,y,z)$ 是一个勾股数三元组，那么 $(ax,ay,az)$ 也是个勾股数三元组。那么，不妨考虑 $f(n)$ 为斜边长为 $n$ 的勾股三元组数量，$g(n)$ 为斜边长为 $n$ 的本原勾股数，那么有

$F(n) = \sum_{i=1}^n f(i) = \sum_{i=1}^n \sum_{d|i} g(d) = \sum_{d=1}^n \sum_{i=1}^{\lfloor n/d \rfloor} g(i) = \sum_{d=1}^n G(\left\lfloor\frac{n}{d}\right\rfloor)$

考虑到所有本原勾股三元组可以表示为 $(i^2-j^2, 2ij, i^2+j^2)$，其中 $\gcd(i,j)=1$，并且不能同时为奇数，那么有

$\def{\isodd}{\in odd} \def{\iseven}{\in even} \begin{aligned} G(n) & = \sum_{i \isodd} \sum_{j \iseven} [i^2+j^2 \le n] [\gcd(i,j)=1] \\ & = \sum_{d=1}^{\sqrt{n/2}} \mu(d) \sum_{i \isodd} \sum_{j \iseven} [i^2+j^2 \le n] [d|i] [d|j] \\ & = \frac{1}{2} \sum_{d=1}^{\sqrt{n/2}} \mu(d) \left( \sum_{i} \sum_{j} - \sum_{i \isodd} \sum_{j \isodd}\right ) [i^2+j^2 \le n] [d|i] [d|j] \\ & = \frac{1}{2} \sum_{d=1}^{\sqrt{n/2}} \mu(d) \left( \sum_{i} \left\lfloor\frac{\sqrt{n-i^2d^2}}{d}\right\rfloor - [d \isodd] \sum_{i \isodd} \left\lfloor\frac{\frac{\sqrt{n-i^2d^2}}{d}+1}{2}\right\rfloor \right ) \end{aligned}$

感觉写到这里一脸复杂度会超掉的样子，那么继续化简试试。

$\def{\isodd}{\in odd} \begin{aligned} F(n) & = \sum_{c=1}^n G(\left\lfloor\frac{n}{c}\right\rfloor) \\ & = \frac{1}{2} \sum_{c=1}^n \sum_{d=1}^{\sqrt{n/2c}} \mu(d) \left( \sum_{i} \left\lfloor\frac{\sqrt{n/c-i^2d^2}}{d}\right\rfloor - [d \isodd] \sum_{i \isodd} \left\lfloor\frac{\frac{\sqrt{n/c-i^2d^2}}{d}+1}{2}\right\rfloor \right ) \\ & = \frac{1}{2} \sum_{d=1}^{\sqrt{n/2}} \mu(d) \sum_{c=1}^n \left( \sum_{i} \left\lfloor\sqrt{\frac{n}{cd^2}-i^2}\right\rfloor - [d \isodd] \sum_{i \isodd} \left\lfloor\frac{1}{2}\left(\sqrt{\frac{n}{cd^2}-i^2}+1\right)\right\rfloor \right ) \end{aligned}$

那么可以发现最后那个对 $c$ 求和的东西对外只和 $n/d^2$ 有关系，对内可以数论分块，一脸时间复杂度 $O(n^{3/4})$ 的模样，复杂度够了，加个记忆化。冲冲冲~

#include <bits/stdc++.h>
using namespace std;
typedef long long lld;

inline int sqrtd(int x)
{
    int tot = sqrtl(x);
    while (tot * tot < x) tot++;
    while (tot * tot > x) tot--;
    return tot;
}

inline pair<lld,lld> c(int n)
{
    static unordered_map<int,pair<lld,lld>> _C;
    if (_C.count(n)) return _C[n];
    lld p1 = 0, p2 = 0; int m = sqrtd(n);

    for (int i = 1, j = m; i <= m; i++)
    {
        while (i * i + j * j > n) j--;
        p1 += j;
        if (i & 1) p2 += (j+1)/2;
    }

    return _C[n] = make_pair(p1, p2);
}

inline lld h(int n, bool odd)
{
    lld ans = 0;

    for (int L = 1, R; L <= n; L = R+1)
    {
        R = n / ( n / L );
        auto cur = c(n/L);
        ans += (R-L+1) * cur.first;
        if (odd) ans -= (R-L+1) * cur.second;
    }

    return ans;
}

int main()
{
    const int MAXN = 1e6+5;
    static int pri[MAXN], pcn;
    static char isnp[MAXN], miu[MAXN];
    miu[1] = 1;

    for (int i = 2; i < MAXN; i++)
    {
        if (!isnp[i]) pri[pcn++] = i, miu[i] = -1;
        for (int j = 0; j < pcn && i * pri[j] < MAXN; j++)
        {
            isnp[i * pri[j]] = 1;
            if (i % pri[j] == 0) break;
            miu[i * pri[j]] = -miu[i];
        }
    }

    int T, n;
    scanf("%d", &T);

    while (T--)
    {
        scanf("%d", &n);
        lld ans = 0;
        for (int d = 1; d * d <= n; d++)
            if (miu[d]) ans += miu[d] * h(n/d/d, d&1);
        printf("%lld\n", ans/2);
    }

    return 0;
}