AtCoder が提供しているABC（AtCoder Beginner Contest）330 のC問題をC++とPythonで解いてみました。ABC330は、2023年11月25日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
C問題 Minimize Abs 2（Difficulty : 519）
問題はリンク先をご覧ください。
2乗の値を計算しておき二分探索です。AtCoder Problems による Difficulty は 519 でした。
解答案
2乗の値を計算しておき、単純に2重ループで計算すると TLE（Time Limit Exceeded）判定となります。$1 \leqq D \leqq 2 \times 10^{12}$ となっているため、このプログラムの計算量は、$\sqrt{D}^2 = D$ となるためです。以下のプログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define INF (1ULL << 60)
int main()
{
ll d;
cin >> d;
vector<ll> a;
ll n = 0;
while (n * n <= d) {
a.push_back(n * n);
++n;
}
a.push_back(n * n);
ll result = INF;
for (int i = 0; i < a.size(); ++i) {
for (int j = 0; j < a.size(); ++j) {
result = min(result, abs(a[i] + a[j] - d));
}
}
cout << result << endl;
return 0;
}このような場合は、二分探索の出番となります。以降、プログラムと合わせて解説します。
C++ プログラム例（ABC330C）
プログラムの補足です。
- $x=0$ の場合、$| y^2\, -\, D |$ の値を最小にする $y^2$ を探すために候補は、$n^2 \leqq D$ となる $n^2$ に加えて、ひとつ大きな $n$ に対しても値を登録しておきます（18行目で処理をしています）。
- $y^2 \geqq D\, -\, x^2$ となる最小の値 $y^2$ を二分探索で求めます。これは式が正の値の場合の最小値となります。もし $y \neq 0$ ならば、$y$ の値を一つ少なくして目的の式の値を求めると、絶対値を取る前が負の値の場合の最小値となります（22－27行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define INF (1ULL << 60)
int main()
{
ll d;
cin >> d;
vector<ll> a;
ll n = 0;
while (n * n <= d) {
a.push_back(n * n);
++n;
}
a.push_back(n * n);
ll result = INF;
for (int i = 0; i < a.size(); ++i) {
auto it = lower_bound(a.begin(), a.end(), d - a[i]);
result = min(result, abs(a[i] + *it - d));
if (it != a.begin()) {
--it;
result = min(result, abs(a[i] + *it - d));
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC330C）
Pythonでは、bisectモジュールをimportすると二分探索が使えます。基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 330 C"""
import bisect
INF = 1 << 60
d = int(input())
a = []
n = 0
while n * n <= d:
a.append(n * n)
n += 1
a.append(n * n)
result = INF
for i in range(len(a)):
j = bisect.bisect_left(a, d - a[i])
result = min(result, abs(a[i] + a[j] - d))
if j != 0:
j -= 1
result = min(result, abs(a[i] + a[j] - d))
print(result)こちらも「AC」と判定されました。
最後に
2乗の値を計算して全検索すると、候補数 $N$ に対して計算量 $O(N^2)$ ですが、二分探索を使うことによって、$O(N\log N)$ にすることができました。定番手法だと考えています。
引き続き ABC の問題を紹介していきます。
