AtCoder Regular Contest 125 - AtCoder
B - Squares题目要求的为\(\sum_{x=1}^n\sum_{y=1}^n[f(x^2-y)]\)，其中\(f(x)\)判断这个数是否是完全平方数，包括零。
我们考虑逆向枚举，设存在\(t\)使得\(x^2-y=t^2\)，那么题目要求的可以转换成\(\sum_{x=1}^n\sum_{t}[1\leq x^2-t^2\leq n]\)
因为\(x^2-t^2\)的极大值为\(x^2\)所以我们的\(x\)实际上最多为\(\sqrt n\)所以上式可以进一步变为\(\sum_{x=1}^{\sqrt{n}}\sum_{t=0}^{\sqrt n}[1\leq x^2-t^2\leq n]\)。
通过打表我们可以发现\(x^2-t^2\)可能的值如下：
\(x=1\)：\([1]\)
\(x=2\)：\([4,3]\)
\(x=3\)：\([9,8,5]\)
\(x=4\)：\([16,15,12,7]\)
\(\cdots\)
然后我们重新将这些值从小到大排列：
\(x=1\)：\([1]\)
\(x=2\)：\([3,4]\)
\(x=3\)：\([5,8,9]\)
\(x=4\)：\([7,12,15,16]\)
\(\cdots\)
观察每一列，发现这是个递增函数，比如第一列\([1,3,5,7]\)，则通项公式为\(a_n=1^2+1\cdot2\cdot(n-1)\)，第二列\([4,8,12]\)则为\(a_n=2^2+2\cdot2\cdot(n-1)\)，那么第\(i\)列就是\(a_n=i^2+i\cdot2\cdot(n-1)\)。
于是我们考虑枚举每一列，并找到最大的\(a_n\leq n\)，这个过程可以二分做到，因为每一列一开始的值为\(i^2\)，所以我们只需要枚举到\(i^2\leq n\)即可，总复杂度大约为\(\sqrt nlog_2n\)。比官方做法要多了个\(log\)。
#include <bits/stdc++.h>
using namespace std;
#define int unsigned long long
const int MOD = 998244353;
int f(int x, int n) {
return x * x + x * 2 * (n - 1);
}
signed main(int argc, char *argv[]) {
ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
int n;
cin >> n;
int res = 0;
auto calc = [&](int x) -> int {
int l = 1, r = 1e12, res = 0;
while (l <= r) {
int mid = (l + r) >> 1;
if (f(x, mid) <= n) res = mid, l = mid + 1;
else r = mid - 1;
}
return res;
};
for (int i = 1; i * i <= n; ++i) {
res = (res + calc(i)) % MOD;
}
cout << res << '\n';
system("pause");
return 0;
}
