本問題のような座標の決め方をしたHEXグリッドにおいて、$(0,0)$ と $(x,y)$ との距離は $\max(|x|,|y|,|x-y|)$ で求めることができるらしい。なるほど。
まずはこれを知ってるか気付くかしないと難しい。
で、各マスを改めて $(x,y,x-y)$ の3次元空間上の点として置きなおした時、 ある $D \times D \times D$ の立方体の中に入る点同士は、どのように選んでも距離が $D$ 以下となる、ということができる。
これなら、立方体の中の点の個数を数えて、それぞれを選ぶ/選ばないを2の何乗かで計算できるので、かなり考えやすくなった。
とはいえ、じゃあどのような立方体を考えればいいのか?
中途半端な位置にある立方体を考える必要は無いので、3軸方向のそれぞれ一方の面には点があると考えてよい。
このように、端に来る点を固定するのはよくある考え方である。
ここでは、各軸方向で小さい座標の方の面に点がある立方体を考える。
ただ、固定する3点を全探索し、さらにそれぞれで立方体の中に入る点を数えると、$O(N^4)$ となり、TLEする。
固定するのはまず2点だけでよく、最後に3点目を固定すると $O(N^3)$ や $O(N^3 \log{N})$ となる。
まず、x座標最小の点 $P=(X_P,Y_P)$ と、y座標最小の点 $Q=(X_Q,Y_Q)$ を固定する。
$P$ と $Q$ が同じ点である場合もあり得る。
$P,Q$ は必ず選ぶようにする前提なので、以下の場合は処理しない。
その上で、x座標が $[X_P,X_P+D]$、y座標が $[Y_Q,Y_Q+D]$ の正方形に含まれ、 かつ $P$ とも $Q$ とも距離が $D$ 以下である点をリストアップする。
D ┼・─────┼ │ ・ │ P ・ ・ │ │ ・ │ ・ │ ┼─Q────┼ D
リストアップされた点は、個別には $P,Q$ とともに選ぶことができるが、 しかし、残る1軸(x-y軸)によっては、リストアップされた点同士は一緒に選べないかも知れない。
ただ、ここまで来れば残りは1次元なので、$x-y$ でソートし、 ここで3軸目の最小座標の点を固定してやれば、どこまで選ぶことができるかは尺取りや二分探索等で簡単にわかる。
固定した3点目を $R$、その $x-y$ の値を $Z_R$ とする。
リストアップされた点のx-yのソート結果例 R Zi: [ -3 -1 -1 1 2 4 ]
$[Z_R,Z_R+D]$ 内に、$R$ を除いてたとえば4点が含まれるとすると、 $\{P,Q,R\}$ を全て含む選び方に対して、その4点は自由に選べる。よってこの3点を固定した場合は $2^4$ 通りとなる。
これを、$R$ を全通り(上記の例なら6点)ずらしながら、結果を合計する。
最後に、$\{P,Q\}$ の2点($P=Q$ の場合は1点)だけしか選ばれないという1通りを入れれば、 $P,Q$ を固定した場合の数え上げが完了する。
x座標やy座標が同じ点が複数あるとき、特に重複がややこしくなる。
主に2つの状況で、重複を防ぐよう注意が必要となる。
x座標最小 $P$ とy座標最小 $Q$ は、同じ点である場合もある。
y │・ │ ・ ② │ ・ ─①─③── x │
「x座標最小を①、y座標最小も①」と固定した場合の処理では、 ②や③も正方形に含まれる点としてリストアップされ、{①,②} や {①,②,③} などが数え上げられる。
ここで「x座標最小を②、y座標最小を①」と固定した場合も処理してしまうと、{①,②} などが重複して数え上げられてしまう。
「x座標最小を①、y座標最小を③」とした場合も同様。
よって、「2点が選ばれたとき、その座標のいずれかが同じ2点」である場合は、 既にx軸最小・y座標最小とも同じ点とした場合で数えられているため、処理をしてはいけない。
y │ ③ ・│ ・ │ ・ ② │ ・ ① ┼─④─⑤───⑥─ x ・
たとえば上記の例で、$(①,②,④,⑤)$ は、
など、計4通りで重複して数えてしまう。
これを防ぐ方法としては、正方形に含まれる点をリストアップするとき、
というように、同じのx座標の点の中では、固定した点がy座標でも最小であるように選べばよい。
とにかく、様々な局面で「最小である1点を固定」という考え方が重宝する問題。