高校数学みたいな問題でした． 脳みそが働いていなかったこともあって，相当な時間詰まってしまいました…
脳みそが死んでいても手を動かすだけでこの手の問題を解けるように整理しておきます．
問題概要
ピタゴラス数とは $$ a^{2}+b^{2}=c^{2} \tag{1} $$ を満たす正の整数 $ (a, b, c) $ の組である． 誤ったピタゴラス数の定義 $$ c=a^{2}-b \tag{2} $$ を考える．$ (3, 4, 5) $ はピタゴラス数正しい定義と誤った定義両方を満たす． このように，両方の定義を満たし $ 1 \leq a \leq b \leq c \leq n $ を満たすような $ (a, b, c) $ の組の数を求めよ． テストケースは $ t $ 個与えられる．
制約
- $ 1 \leq t \leq 10^{4} $
- $ 1 \leq N \leq 10^{9} $
考察
式 $ (1) $ と $ (2) $ を両方満たさなければならないため，連立させる． 共通する $ a^{2} $ を排除したいので辺々足し合わせて整理する．
上の式 $ (3) $ を $ (2) $ 代入して（別に $ (1) $ でも良い）
$ a $ と $ c $ の関係式が得られた．組の数え上げの典型であるところの関係式が得られたら文字を固定するをやる． こういうのは次数が高い文字から決めうつのが好ましい．（ $ a $ と $ c $ の上限が同じなので次数が高い $ a $ の方が探索範囲が狭まる）
式 $ (4) $ と制約 $ c \leq n $ より $ a^{2} \leq 2n-1 $ を満たす $ a $ を考えればいいことが分かる． また 式 $ (4) $ から $ c $ が整数である条件は $ a $ が奇数であれば良いことがわかる． $ b=c-1 $ なので $ c $ が整数であれば $ b $ も整数なのでこちらも同時に条件を確認できる．
最後に $ 1 \leq a \leq b \leq c $ の条件を確認する．
- $ a $ を $ 1 $ からインクリメンタルに全探索するので $ 1 \leq a $ は満たされる．
- $ b = \frac{a^{2}-1}{2} $ と計算できるので $ a \leq b $ はそのまま条件確認すれば良い．
- $ b \leq c $ は常に満たされている．
$ a^{2} $ は $ a>0 $ で単調増加するので $ 1 $ つのテストケースを $ O( \sqrt{n} ) $ で計算ができる． $ O(T \sqrt{n} ) $ は正直厳しいケースがありそう． $ n $ の最大値を $ n _ {max} $ として $ a^{2} \leq 2 \cdot n _ {max} - 1 $ までについて前計算でテーブルをすれば，二分探索によって各テストケースに $ O( \log{n} ) $ で答えられる． こうすることで，時間計算量 $ O( \sqrt{n_{max}} + T \log{n}) $ で安心して Submit することができる．
別解
計算量はそう変わらないんですが，別の解法を一応示します． $ a^{2} $ が $ a>0 $ で単調増加で $ b=c-1 $ とおけば，$ a \leq b \leq c $ も $ a=3 $ 以上の時は満たされている（ $ f(a)=b-a=\frac{a^{2}-1}{2}-a $ のグラフを思い浮かべれば自明かと思います）．
つまり，$ 2n-1 $ 以下で $ 3 $ 以上の奇数の平方数を数え上げれば問題文の条件を全て満たす． これは $ a=2m+1 $ などと表現すれば二分探索で求めることができる． 前計算が要らずに $ O(T \log{n}) $ になり，うれしい（？）
謝辞
前計算をしない $ O(T \sqrt{n} )$ が普通に 358 ms で通りました．実行時間の見積もりが下手すぎて涙…。
個人的ポイント
- 連立方程式が出てきたらとりあえず文字を減らすよう頑張る．
- 文字が $ 2 $ つ以下にできたら後は大体算数か二分探索．
- 次数が大きい方から決め打つと，十中八九嬉しい．
- これはもう脳死でやれるだろ（自戒）
- 整数組数え上げ問題はちゃんと最初の条件をちゃんと確認する．
- ちゃんと確認しないと頭悪い日だと混乱してしまう…
- 整数条件と不等式
- なんか軽い $ 10^{8} $ ループぐらいならもう脳死で投げちまうか．
- C++ は最強なので大体通る気がしてきたぞ．
