ABC397 D - Cubes
https://atcoder.jp/contests/abc397/tasks/abc397_d
先日のD問題が運良く解けて嬉しかったので記事にします。
前半は公式解説と同じですが後半が違います。嘘解法だったらすみません。
追記：
行間を飛ばしているところが多々あると思ったのでもう少し丁寧に書きました。間違っていそうなところも直しました。
コンテスト中の自分の動き
まず考えたのは公式解説の O(√N) 解法です。10^6 までの数を3乗したものを計算して配列として持っておきます。その配列の中身を見て差が k^3 となる k があればOKです。しかしこれは大きな勘違いです。立法数同士の差が最大 10^18 になるのでもっと大きな数値まで計算する必要がありました。これを提出したせいでペナルティ1です。余計なprint文を入れていたのでその前にもペナルティ1、こちらはもったいなかったです。
考え直しました。「こういうときは数学」というのがよくあるパターンです。因数分解できるのはすぐにわかったので
\begin{align}
N &= x^3 - y^3 \\
&= (x - y)(x^2 + xy + y^2)
\end{align}
とします。中学生の頃を思い出しますね。N が 2 つの整数の積で表せるようになりますから、これならなんとかいけそうです。この方針で考えを進めていき、35分ほどで問題を解くことができました。
考察
N を 2つの整数 P, Q の積と考えます。P, Q は N の約数ですが、最大 10^18 の N の約数を調べるには多大な計算量を要するのでなるべく探索範囲を小さくしたいところです。そこでPとQの大小関係を調べます。小さい方だけを探索してそれが小さい計算量に収まれば嬉しいですよね。
\begin{align}
P &= x - y\\
Q &= x^2 + xy + y^2
\end{align}
とします。まあ普通に考えれば P < Q ですね。ここで中学生の頃に習った公式
P^2 = (x-y)^2 = x^2 - 2xy + y^2
を思い浮かべます。これと先ほどの Q の式、似てるなあと昔から思っていました。これらを見比べると
Q - P^2 = 3xy > 0
であることがわかります。つまり
\begin{align}
P^2 &< Q\\
\end{align}
ですね。PもQも正なので
\begin{align}
P &< Q\\
\end{align}
でもあります。P^2 < Q を利用して N の式を変形します。
\begin{align}
P Q &= N \leq 10^{18}\\
P P^2 &< P Q \leq 10^{18}\\
P^3 &< 10^{18}\\
P &< 10^6
\end{align}
これにより N を 1 から 10^6 までの整数で割ることで全ての (P, Q) を列挙することができ、計算も間に合います。
ここまでの情報から
\begin{align}
x - y &= P\\
x y &= \frac{Q - P^2}{3}
\end{align}
ということがわかります。ここまで来れば答えにたどり着けそうです。ここで、P^2やQの式と似た形の式を思い浮かべます。基本的な公式です。
(x+y)^2 = x^2 + 2xy + y^2
です。これは
\begin{align}
x^2 + 2xy + y^2 &= x^2 - 2xy + y^2 + 4xy\\
&= (x-y)^2 + 4xy
\end{align}
と変形できますから、x-y と xy を利用して x+y が求められます。あとは簡単に x と y が求められますね。
\begin{align}
x &= \frac{(x+y)+(x-y)}{2}\\
y &= \frac{(x+y)-(x-y)}{2}
\end{align}
です。和差算です。
というわけで
- 1 <= P <= 10^6 なる P に対し、P が N の約数なら (P, Q) を答えの候補とする。
- Q - P^2 = 3xy なので、Q - P^2 が 3 の倍数になる場合は 3 で割って xy を得る。
- P^2 に 4xy を足して (x+y)^2 を得る。この数字が平方数なら x+y が得られる。
という流れになります。x, y が整数であることを利用して、条件に合わない候補を減らしていきます。
平方数の判定
最後の「この数字が平方数なら」の判定時に悩んでしまいました。math.sqrt() を使うといかにも誤差が出そうです……誤差についてはもっともっと正確に知っておく必要があると思うのですが今のところ「なんとなく怪しい」ぐらいの腫れ物扱いです。そこで平方数をあらかじめ挙げておき、その中に (x+y)^2 に該当する値があれば平方数と判定することにします。
ところが……焦っておりどのぐらい大きな値まで考えればいいのかよくわかりませんでした。とりあえず 10^6ぐらいでいいか、もっと余裕を見るかと思って 10^7 までの数字を 2 乗したものを用意しました。そして提出。……TLEになりました。コードテストで試してみたところ、この配列を用意するだけで TLE になってしまいます。何事も焦ってはいけませんね。続いて 2*10^6 までに修正して提出してみましたが、WA が出ました。もっと大きな数まで見ないといけないようです。
提出を繰り返しながら二分探索して上限を探るわけにもいきませんから、先ほど怪しんだ sqrt を使ってみます。Google検索してみると整数の判定には is_integer() というメソッドがあるようなのでそれを活用しました。すると、無事 AC を得ることができました。これなら先に一か八かの賭けで sqrt を使ったコードを提出してみればよかったかもしれません。
というわけで、最後が若干怪しいものの AC することができました。after_contest で弾かれるようになったらすみません。
後から考えてみたのですが、本当に平方数かどうかが疑わしいのなら検算すればいいかもしれません。得られた値を2乗して下の値と一致すれば正解です。不安なら得られた値の±1ぐらいまで調べてみてもいいかもしれません。
実装
import math
N = int(input())
ans = "-1"
for p in range(1, min(N+1, 10**6+1)):
if N % p != 0:
continue
q = N // p
if p**2 >= q:
continue
if (q - p**2) % 3 != 0:
continue
x_multi_y = (q - p**2) // 3
x_plus_y_2 = p**2 + 4 * x_multi_y
x_plus_y = math.sqrt(x_plus_y_2)
if not x_plus_y.is_integer():
continue
x_minus_y = p
x = (x_plus_y + x_minus_y) // 2
y = (x_plus_y - x_minus_y) // 2
ans = str(int(x)) + " " + str(int(y))
break
print(ans)