AtCoder Begginer Contest 389 D

はじめに

先日開催されたAtCoder Beginner Contest 389のD問題を解きました。AtCoderの問題を解くために必要な考え方や基本的なアルゴリズムを多く学べたので、記録します。

問題概要

一辺の長さが$1$である正方形が平面の上に敷き詰められています。ある正方形の中心を中心とする半径$R$の円に含まれる正方形はいくつあるか求めてください。

制約

• $1\leq R\leq 10^{6}$
• $R$は整数

考え方

対称性を考えれば、第一象限に含まれる正方形の数を数えて、$4$ 倍すれば良さそうです。ただし、軸の上に乗っているものは先に数えておく必要があります。$(i,j)$を中心とする正方形を$s_{i,j}$と表します。次の軸上の正方形は円に含まれます¹。

• $s_{0,0}$
• $s_{0,\pm j}\quad(1\leq j\leq R-1)$
• $s_{\pm i,0}\quad (1\leq i \leq R-1)$

変数は一つずつ動かす

次に、第一象限のうち、半径$R$の円の内部である領域$D$に含まれるものの数を求めます。今動かすことのできる変数として、正方形を指定する$i,j$が与えられています。変数が2つ以上ある場合は、ひとつずつ動かして考えるのが定石のようです。そこで、初めに$i$を$1$ に固定して$D$に含まれる正方形の数$N(1)$を数え、次に$i$を$2$に固定して$N(2)$を数える、という操作を$i=R-1$となるまで繰り返します。この操作ののち、$\sum_{i=1}^{R-1}N(i)$が$D$に含まれる正方形の数となります。

条件を言い換える

このとき、正方形$s_{i,j}$が$D$に含まれる必要十分条件は
$$ \left(i+\frac{1}{2}\right)^2 +\left(j+\frac{1}{2}\right)^2 \leq R^2$$
です。この条件が満たされているとき、$s_{i,j}$は$D$に含まれます。逆に、この条件が成り立たないとき、$s_{i,j}$は$D$ に含まれません。

列を二分する条件に着目する

さらに重要な性質として、$i$列目の正方形の列$s_{i,1},s_{i,2},\dots$を下から順番に見ていったとき、ある正方形$s_{i,j}$が存在し、$s_{i,j}$以下にある正方形は$D$に含まれ、$s_{i,j}$より上にある正方形は$D$に含まれません。つまり、この境界は$i$列目の正方形の列を二分します。

計算量を見積もる

では、$N(i)$を求める具体的な方法を考えましょう。次のような単純な探索を考えます。$j$を$R-1$から$1$ずつ減らしていって、$D$に含まれるようになったところで、繰り返しをやめます。そのときの$j$を$j_{\mathrm{m}}(i)$とすると、$j_{\mathrm{m}}(i)$が$N(i)$に他なりません。このような探索の場合、繰り返しの対象となる正方形の数が大雑把に見積もって$R^2$のオーダーであることから、制限時間を超過してしまいます。一応、数式でも確認しておきます。必要な繰り返し回数は

\begin{align}
(繰り返し回数)&=\sum_{i=1}^{R-1}\left[R-1-\mathrm{floor}\left(\sqrt{R^2 - \left(i+\frac{1}{2}\right)^2}-\frac{1}{2}\right)\right]\\
& >(R-1)^2 - \sum_{i=1}^{R-1}\left(\sqrt{R^2 - \left(i+\frac{1}{2}\right)^2}-\frac{1}{2}\right)
\end{align}

$\sum_{i=1}^{R-1}\sqrt{R^2 - \left(i+1/2\right)^2}$の部分は$\pi R^2/4$より小さいですから

\begin{align}
(繰り返し回数)&>(R-1)^2 - \left(\frac{\pi R^2}{4} - \frac{1}{2}(R-1)\right)\\
& = \left(1-\frac{\pi}{4}\right)R^2 -\frac{3}{2}R +\frac{1}{2}
\end{align}

となります。したがって、この探索の計算量は常に少なくとも$O(R^2)$です。制約を考えると、実行時間を超過します。そこで、計算方法を工夫する必要があります。

二分探索

二分探索では、以下のような状況で目的の値を効率よく探索することができます。

• ある条件が列を二分する
• その条件を満たす最大または最小のインデックスを知りたい
  今の問題設定では、各$i$について
  $$
  \left(i+\frac{1}{2}\right)^2 +\left(j+\frac{1}{2}\right)^2 \leq R^2
  $$
  を満たす最大の$j(=j_{\mathrm{m}}(i))$を求めることができます。計算量は多く見積もっても$O(R\log(R))$です。

def is_included(i, j):
	'''
	i, jがDに含まれているか確かめる
	'''
	return (2 * i + 1) * (2 * i + 1) + (2 * j + 1) * (2 * j + 1) <= 4 * R * R

R = int(input())

count = 0
for i in range(1, R):
	#はじめ、jの最大値の候補は[0,R)
	top = R
	bottom = 0
	mid = (top + bottom) // 2
	
	while top - bottom > 1:
		if is_included(i, mid):
			bottom = mid
		else:
			top = mid
		mid = (top + bottom) // 2
	
	count += mid

count *= 4
count += 4 * (R-1) + 1
print(count)

尺取り法

$j_{\mathrm{m}}(i)$は$i$について単調減少します。このことに着目すると、初めに述べた探索を改善することができます。単調減少することから、繰り返しを減らすことができるイメージです。

def is_included(i, j):
	'''
	i, jがDに含まれているか確かめる
	'''
	return (2 * i + 1) * (2 * i + 1) + (2 * j + 1) * (2 * j + 1) <= 4 * R * R

R = int(input())

total = 4 * (R-1) + 1

j = R - 1
count = 0

for i in range(1, R):
	while !is_included(i, j):
		j -= 1
	count += j

total += 4 * count

print(total)

まとめ

• 変数は一度に全て動かさないで、一つずつ動かす
• 問題文で与えられた条件を、分かりやすい必要十分条件に言い換える
• 列を二分する条件に着目する(二分探索)
• 単調性に着目する(尺取り法)

1. 図を書いてみれば分かるように、これらの正方形が円に含まれる条件は、　$(1/2, R - 1/2)$と原点の距離が$R$以下であることですが、これは確かに成り立ちます。 ↩