P14139 题解

有意思的题目。

首先把式子拆成两部分：

\[\begin{aligned} A &= \sum_{x = 1}^n \left \lfloor \dfrac{x^2}{n} \right \rfloor \\ B &= \sum_{x = 1}^n \left \lfloor \sqrt{nx} \right \rfloor \end{aligned} \]

个人认为 \(B\) 更好算，所以先算 \(B\)。

考虑 \(\left \lfloor \sqrt{x} \right \rfloor\) 的意义，注意到：

\[\left \lfloor \sqrt{x} \right \rfloor = \sum_{y^2 \le x} 1 \]

则：

\[\begin{aligned} B &= \sum_{x = 1}^n \left \lfloor \sqrt{nx} \right \rfloor \\ &= \sum_{x = 1}^n \sum_{y^2 \le nx} 1 \\ &= \sum_{y = 1}^n \sum_{x = \lceil y^2 / n \rceil}^n 1 \\ &= \sum_{y = 1}^n \left( n - \left \lceil \dfrac{y^2}{n} \right \rceil + 1 \right) \\ &= n^2 + n - \sum_{y = 1}^n \left \lceil \dfrac{y^2}{n} \right \rceil \end{aligned} \]

我们发现，这个 \(\sum\) 和 \(A\) 很相似啊，所以：

\[\begin{aligned} A + B &= n^2 + n + \sum_{x = 1}^n \left( \left \lfloor \dfrac{x^2}{n} \right \rfloor - \left \lceil \dfrac{x^2}{n} \right \rceil \right) \\ &= n^2 + n - \sum_{x = 1}^n \left[ n \nmid x^2 \right] \\ &= n^2 + \sum_{x = 1}^n \left[ n \mid x^2 \right] \end{aligned} \]

到这里式子就推不动了，所以只能从 \(n \mid x^2\) 入手。

设 \(n = ab^2\)，其中 \(b\) 为最大的满足 \(b^2 \mid n\) 的整数。则 \(n \mid x^2 \iff ab^2 \mid x^2 \iff a^2 b^2 \mid x^2 \iff ab \mid x\)。那么：

\[\begin{aligned} A + B &= n^2 + \sum_{x = 1}^{ab^2} \left[ ab \mid x \right] \\ &= n^2 + \left \lfloor \dfrac{ab^2}{ab} \right \rfloor \\ &= n^2 + b \end{aligned} \]

那么只需要算 \(b\) 就可以了。受部分分启发，我们先做 \(n\) 为完全平方数的情况。容易发现，此时 \(b = \sqrt{n}\)。

那么若 \(n\) 不为完全平方数呢？如果 \(n\) 只有两个质因子，那么 \(b = 1\)。否则每个质因子都 \(\le 10^6\)，可以预处理 \(\le 10^6\) 的质数，然后对 \(n\) 分解质因数。设 \(n = \prod \limits_i p_i^{\alpha_i}\)，则 \(b = \prod \limits_i p_i^{\lfloor \alpha_i / 2 \rfloor}\)。

时间复杂度：\(O \left( \dfrac{T \sqrt[3]{n}}{\ln(n)} \right)\)。

#include <bitset>
#include <cmath>
#include <iostream>
#include <vector>

using namespace std;
using i64 = long long;

constexpr int MOD = 1e9 + 7, CURT_N = 1e6;

vector<int> primes;

void precompute()
{
    bitset<CURT_N + 1> not_prime(3);

    for (int i = 2; i <= CURT_N; ++i) {
        if (not_prime.test(i))
            continue;

        primes.push_back(i);

        for (i64 j = i64(i) * i; j <= CURT_N; j += i)
            not_prime.set(j);
    }
}

inline int add_mod(i64 x, i64 y) { return (x % MOD + y % MOD + MOD) % MOD; }
inline int mul_mod(i64 x, i64 y) { return (x % MOD) * (y % MOD) % MOD; }

inline bool is_square(i64 x)
{
    i64 root = sqrt(x);
    return root * root == x;
}

int pow_mod(i64 x, i64 y)
{
    int res = 1;

    for (; y > 0; y >>= 1, x = mul_mod(x, x)) {
        if (y & 1)
            res = mul_mod(res, x);
    }

    return res;
}

struct Solution {
    void main()
    {
        i64 n;
        cin >> n;

        int sq = 1;
        i64 orig_n = n;

        for (auto p : primes) {
            if (n % p != 0)
                continue;

            int cnt = 0;

            while (n % p == 0) {
                n /= p;
                ++cnt;
            }

            sq = mul_mod(sq, pow_mod(p, cnt >> 1));
        }

        if (is_square(n))
            sq = mul_mod(sq, sqrt(n));

        cout << add_mod(sq, mul_mod(orig_n, orig_n)) << '\n';
    }
};

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int t;
    cin >> t;
    precompute();

    while (t-- > 0)
        Solution().main();

    return 0;
}