ARC125 B - Squares - ichinoseac’s diary

https://atcoder.jp/contests/arc125/tasks/arc125_b

自分の解法メモです。

問題

整数 $N$ が与えられる．以下の条件を満たす整数の組 $(x, y)$ の個数を998244353で割った余りを求めよ．

$x^ 2 - y$ は平方数（0も平方数とする）

制約

$1 \leq x, y \leq N$
$1 \leq N \leq 10^ {12}$

考察

以下、 $x, y, N$ は制約を満たすものとします。

$∃x,y [ x^ 2 - yが平方数 ] \\ ⇔∃x,y,k [ 0 < k \leq x かつ x^ 2 - y = (x-k)^ 2 ] （以下 k は整数） \\ ⇔∃x,y,k [ 0 < k \leq x かつ x^ 2 - (x-k)^ 2 = y ] \\ ⇔∃x,k [ 0 < k \leq x かつ 1 \leq x^ 2 - (x-k)^ 2 \leq N ] ( \because 1 \leq y \leq N ) \\ ⇔∃x,k [ 0 < k \leq x かつ \frac{k^ 2 + 1}{2k} \leq x \leq \frac{k^ 2 + N}{2k} ] \\ ⇔∃x,k [ 0 < k \leq x \leq \frac{k^ 2 + N}{2k} ] ( \because k = \frac{k^ 2 + k^ 2}{2k} \geq \frac{k^ 2 + 1}{2k} )$

$k$ を固定したとき、これを満たす整数 $x$ の個数は $k$ から $\lfloor \frac{k^ 2 + N}{2k} \rfloor$ までの $\lfloor \frac{k^ 2 + N}{2k} \rfloor - k + 1$ 個。

$k$ は $k \leq \frac{k^ 2 + N}{2k}$ を満たす必要があるので、 $k$ が動く範囲は $k=1, \ldots , \lfloor \sqrt{N} \rfloor$ 。

コード概要

long long total = 0;
for(long long k = 1; k * k <= N; ++k) {
  total += (k * k + N) / (2 * k) - k + 1;
  total %= MOD;
} 
cout << total << endl;

https://atcoder.jp/contests/arc125/submissions/25288020