ABC246 D問題 2-variable function を考える

公式ページ

記事の内容

この問題の公式解説を読んだものの、尺取り法による高速化がすぐに理解できなかったので、考え方をメモ
ついでにQiita初投稿です。

問題の要約

$N \leqq X=a^3+a^2b + ab^2+b^3$ を満たすような非負整数 $(a,b)$ の組の内、
$X$ 最小となる $(a,b)$ を求める。
$N$は標準入力から与えられ、$0 \leqq N \leqq 10^{18}$ を満たす

基本的な考え方

$X=a^3+a^2b + ab^2+b^3$ であるから、$b=0$ としたときに$X=a^3$なので、$a \leqq 10^6$
　⇒ 探索範囲の終点。bも同様。

$a,b$ は非負なので、$X \geqq \ 0$。
$a,b$ は大小関係が無く区別する必要はないので、
$a: 0 \rightarrow 10^6, \ b: 10^6 \rightarrow 0$ へ探索。

これだけでもAC可能。

尺取りの考え方

$b$ の探索範囲を限定する。
最小となる $X$ を探すことに着目。

それまでに探索した中で最小の $X$ よりも、
明らかに大きい $X$ が出てくる $(a,b)$ の組は探索しないように、探索範囲を限定していく。

やり方

① $a,b$ に大小関係が無いから、 $a: 0 \rightarrow 10^6, \ b: 10^6 \rightarrow 0$へ探索するとき、
　$a > b$ となったら終了して良い
② $a=0固定,b:10^6 \rightarrow 0$ へ探索して $N \leqq X$ となる最小の $b_{min(a=0)}$ があった場合、
　$a=1,2,...$ を探索するとき $b_{min(a=0)} \leqq b \leqq 10^6$ の範囲は探索しなくて良い
　 なぜなら、上記範囲では必ず $X>X_{(a,b_{min(a=0)})}$ となるから。

という風にすると、$a=0,1,2,...$ と進むにつれて、
$b$ の探索範囲はどんどん限定されていくことになる。

3次元的に観る

$a=0\sim 10,\ b=0\sim 10$ のZ軸 $X$ のプロット
灰色が $X=2048$ の平面

直感的に、右下（$a=10,b=10$ の辺り）は探索不要だと分かる

ループが進むと灰色の平面が上から降りて来るが、$a,b$ が整数だから離散的な表現しかできず、
引っ掛かって止まるイメージ。

灰色平面とカラーマップ平面の、Z軸方向でより0に近い新しい交点が分かったら、
交点から右下は探索範囲から外される。

ソースコード

$X=2048$ となる $N$ を入力すると上図と同じプロットが得られます。

# Python 3.7.3
def Xcal(a,b):
    return a**3 + a*a*b + b*b*a + b**3

N = int(input())

#----- 基本的な解答 -----#
ans = float('inf')
b = 10**6

for a in range(10**6):
    while Xcal(a,b) >= N and b >= 0:
        ans = min(ans,float(Xcal(a,b)))
        b -= 1
print(f'基本: {int(ans)}')

#----- 尺取りの解答 -----#
ans = float('inf')
a = 0
b = 10**6

while a <= b:
    if Xcal(a,b)>=N :
        ans = min(ans,float(Xcal(a,b)))
        b -= 1
    else:
        a += 1
print(f'尺取: {int(ans)}')

#----- プロットを眺める用のコード -----#
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d, Axes3D
import numpy as np
fig = plt.figure()
ax = fig.add_subplot(projection='3d')

# 3Dプロット
x_data = np.linspace(0,10,100)
y_data = np.linspace(0,10,100)
x, y = np.meshgrid(x_data,y_data)
z = Xcal(x,y)
ax.plot_surface(x,y,z,cmap='jet')

# 解答の平面プロット
x_data = np.linspace(0,10,100)
y_data = np.linspace(0,10,100)
x, y = np.meshgrid(x_data,y_data)
z = (x+y)*0.0 + ans
ax.plot_surface(x,y,z,color='gray')

plt.show()

おわりに

などと書きつつ、当時は良く分からない別法に走ってTLEだったので、次回は頑張ります。
久しぶりにTeX記法？を使って懐かしい気持ちになりました。