https://kntychance.hatenablog.jp/entry/2024/12/15/212627

はじめに

問題の解説をしつつ、競プロ楽しいよ～という布教をする回です。

特に、競プロを未経験の方に、競プロにおいて「典型*1を組み合わせて問題を解く」とはどのようなものなのかを実感してもらえたらと思います。

解説する問題はこいつ ↓ です。とりあえず読んでみてください。

※ 問題文中にあるマンハッタン距離とは、縦横でのみ移動できるとしたときの2点間の距離のことです。厳密には、2点 $(x_1, y_1), (x_2, y_2)$ 間のマンハッタン距離は $|x_1 - x_2|+|y_1 - y_2|$ と表されます*2。

一見するとかなりシンプルな問題ですが、実は出題されているコンテスト（AtCoder Beginner Contest）の中ではかなりの難問に分類されており、典型知識を多重に組み合わせて解く必要があります。

どういった点が難しいのか、どのように解いていくのかを、競プロ関連の前提知識は使わないように（無理なところは適宜端折って）解説していくので、是非最後まで読んで「解けた！」が届く快感を味わってください。

また、できる限り丁寧に解説したつもりではありますが、問題自体が難しいので「解けた！」が万人には届けられないかもしれません。それでも、途中まで読んだ中で「なるほど！」が届けばなと思います。

何が難しいのか

改めて、問題の概要を確認します。

問題概要

二次元座標平面上に $N$ 個の相異なる点が与えられる。
個の質問に答える。各質問の内容は次の通り：
- マンハッタン距離を用いるとして、座標 $(X_i, Y_i)$ から $K_i$ 番目に近い点までの距離は？

実行時間制限: 7sec
$1 \le N \le 10^5$
$1 \le Q \le 10^5$
与えられる各座標の x, y はいずれも $0$ 以上 $10^5$ 以下

この問題、答えを求めること自体はそう難しくはありません。

各質問 $i$ について、以下のような手順で質問に答えることができます：

$\text{dist}[v] := 点 v から (X_i , Y_i) までのマンハッタン距離$ とする配列 $\text{dist}$ を作成する。
$\text{dist}$ を昇順にソートする。
$\text{dist}$ の $K_i$ 番目が答え。

しかし、残念ながらこの解法で実装したコードを提出しても、TLE（Time Limit Exceeded / 実行時間制限超過）となり、不正解となります。プログラムの実行時間が長すぎるのでダメ！と怒られるわけですね。

この問題の実行時間制限は 7sec と設定されています。対して、プログラムの実行時間はどのくらいになるのでしょうか。

上記解法においてボトルネック（= 最も実行時間への影響が大きい箇所）となっているのは 2. のソート部分で、計算量は $\Theta(N \text{log} N)$ である（≒ $N \text{log} N$ の定数倍程度のステップ数を要する）ことが知られています*3。さらに、 $Q$ 回質問することを踏まえると、全体の計算量は $\Theta (QN \text{log} N)$ になります。

一般的なPCが1秒間に処理できるステップ数はおよそ $10^9$ 回程度と言われています。つまり、実行時間はざ～っくりと見積もると

\begin{align} \frac{10^5 \times 10^5 \times \text{log} 10^5}{10^9} \gt 50 \end{align}

という感じで、ガバガバ議論にはなりますが 7sec には到底収まりそうにないことが分かります。

更に言うと、上記解法の 1. だけ見ても計算量は $\Theta (NQ)$ で、この時点で既に 7sec は厳しいです。つまりは、「各質問ごとに各点を見る」という当たり前の処理すらさせてもらえない時間制限になっているのです。

シンプルな見た目をしていますが、実はかなりの無茶ぶりを要求されているということがご理解いただけたでしょうか？

いよいよここからは、競プロの典型知識を用いてこの問題を解きほぐしていきます！

典型1: K番目の取得は決め打ち二分探索

早速ですが、問題の言い換えを行います。

$i$ 番目の質問

座標 $(X_i, Y_i)$ から $K_i$ 番目に近い点までの距離は？

は、次と同値です。

「 $(X_i, Y_i)$ からの距離が $d$ 未満であるような点の個数は $K_i$ 個未満」を満たす、最大の $d$ は？

ちょっと長いので、判定問題として

\begin{equation} P_i (d) :=(X_i, Y_i) からのマンハッタン距離が d 未満であるような点の個数は K_i 個未満であるか？\end{equation}

とおいてあげると、最終的に求めたいものは $P_i (d) = True$ となる最大の $d$ となります。

さて、ここからが重要です。

$P_i (d)$ は、 $d$ についての単調性、すなわち

$d_1 \lt d_2$ なる $d_1, d_2$ について、 $P_i (d_1) = False \Rightarrow P_i (d_2) = False$

という性質を持ちます。（ $d$ を大きく取れば距離 $d$ 未満の範囲にある点の個数は単調増加するので、当然です。）

この性質を用いれば、次のスライドのように範囲を半分ずつ絞っていくことで、判定問題 $P_i (d)$ を高々 $\text{log}_2 (dの範囲)$ 回程度解くだけで所望の $d$ を求めることができます。いわゆる二分探索というやつです。

なお、 $0 \le d \le 2 \times 10^5$ と評価できるので、 $P_i (d)$ を解く回数は高々 $\text{log}_2 (2 \times 10^5) \lt 20$ 回程度です。

このように、何かしらの数値を求める問題に対して単調性のある判定問題に置き換えて二分探索を適用する手法を、決め打ち二分探索と呼びます。

とはいえ、まだ問題はほとんど解決していません。

というのも、決め打った $d$ に対して判定問題 $P_i (d)$ が高速に解けないようでは、判定問題に置き換えた意味がないからです。

ここからは、この判定問題を解く方法を掘り下げていきます。

典型2: マンハッタン距離は45°回転

$(X_i, Y_i)$ からマンハッタン距離が $d$ 未満である領域を図示すると、次のような菱形になります。 $P_i (d)$ は、この領域に含まれる点が $K_i$ 個未満かどうかを判定する問題に相当します。

$(X_i, Y_i)$ からのマンハッタン距離が $d$ 未満の領域（境界線は含めない）

ただ、この菱形というのは色々と扱いづらいです。

そこで、以下のような座標変換を行ってみます：

\begin{align} u &:= x+y \\ v &:= x-y \end{align}

各位置ベクトルに対して45°の回転行列をかけたあと $\sqrt{2}$ 倍すると理解してもよいでしょう。この変換により先ほどの領域は、u-v座標系で表すと次の通り正方形になります。

これにより、正方形区画に点が何個含まれるかを求める問題に帰着することができました。これが解けるかどうかは別として、少なくとも菱形で考えるよりはなんとなく見通しがよさそうですよね。

このように、マンハッタン距離に関する問題はとりあえず45°回転してから考えるというのが典型となっています。

以降はこのu-v座標系で議論を進めていきます。

なお、座標変換により

$0 \le u \le 2 \times 10^5$
$-10^5 \le v \le 10^5$

となっています。

典型3: 平面走査

「正方形区画内の点の個数」を、もう少しだけ扱いやすい形で表現します。

座標平面上の各格子点を二次元グリッドのマス目に置き換え、点が存在するなら 1, しないなら 0 を書き込みます。これにより、二次元グリッド上で正方形区画内の総和を求める問題として定式化できます。

なお、 $u, v$ の範囲を思い出すと、グリッドサイズが縦横いずれも $2×10^5$ となっていることに注意しましょう。（つまり、プログラム上で明にこのグリッドを二次元配列で持つことはできません*4。悲しい。）

さて、このような大きなグリッドに関する問題に対しては、平面走査と呼ばれる典型的なアプローチがあります。まずは、次のスライドをご覧ください。

何がしたいのかは一旦置いておくとして、行ごとに走査しながら配列に値を加算していく様子が伺えます。このように、何かしらのデータを持ち、行（あるいは列）ごとにデータを更新しながら平面を走査することを、総称して平面走査と呼びます。

急にわけのわからない操作をし始めたぞコイツ…？と思われているかもしれませんが、実はこの操作に超重要な秘密が隠されています！

↓この緑部分の総和（4）から…

↓この黄色部分の総和（1）を引くと…

求めたかった正方形区画の総和（4-1 = 3）が得られます！

この理由は単純で、矩形和の引き算に対応しているからです。

ということで、正方形区画の総和は次の要領で求められます：

平面走査を行う。その途中で、次の二つの値を別途保存しておく。
- 黄色部分の総和
- 緑部分の総和
上記で保存した2数の差が答え。

計算量を考えてみましょう。グリッドの横幅を $W$ , $d の範囲$ を $D$ とおくと、

平面走査については、長さの配列を用意したうえで更新が回発生するので、
- 黄色部分、緑部分の総和を求めるのは $\Theta (d)$ で、これは平面走査の $W$ があるためボトルネックにならない。
→ 判定問題 $P_i(d)$ 1回あたり $\Theta (W+N)$
→ 全体 $\Theta (Q(W+N) \text{log} D)$

あれ…？？？？？？全然よくなってない…（´・ω・｀）

典型4: 並列二分探索

平面走査などというトンデモびっくりアプローチまで取り出して頑張ったのに、計算量は何も改善されていませんでした。果たして、今までの努力は何だったのでしょうか…。

ここで、全く新しい視点を導入します。

今までは質問1つごとにどう解くか？ということだけを考えていましたが、思い切って $Q$ 個の質問を一斉に解くことを考えてみましょう。

より正確には、 $Q$ 個の質問に対して二分探索を同時に実施し、 $Q$ 個の判定問題 $P_i (d_i) ~ (1 \le i \le Q)$ を1回の平面走査中にまとめて解くことにします。

今までの議論と同様に考えると、グリッド上に $Q$ 個の正方形区画が与えられ、それぞれの総和を求める問題を解くことになります。実はこれ、正方形が何個だろうとやることは全く変りません。すなわち、平面走査の途中で各正方形区画ごとに対応した黄色部分の総和と緑部分の総和を保存し、最後にそれぞれ差をとればよいです。

ありがとう、平面走査…！！！！

計算量を考えます。

平面走査については、長さの配列を用意したうえで更新が回発生するので、
- 黄色部分、緑部分の総和を求めるのは各 $i$ ごとに $\Theta (d_i)$ で、まとめると $\Theta (\color{red}{QD})$
→ 判定問題 $P_i(d_i) ~ (1 \le i \le Q)$ をまとめて解く処理1回あたり $\Theta (W+N+\color{red}{QD})$
→ 全体 $\Theta ( (W+N+\color{red}{QD}) \text{log} D)$

惜しい！！

ここにきて、配列の区間和を計算する箇所がボトルネックになってしまいました。

典型5: Segment Tree

いよいよ最後の典型です。

平面走査中に行いたい処理を改めて書き下すと、次の2つです：

1点更新
区間和取得

平面走査のデータを通常の配列で持ってしまうと、どうしても区間和取得に区間長分の計算量がかかってしまいます。

そこで登場するのが、競プロにおいて典型中の典型、キング・オブ・典型データ構造である Segment Tree です*5。

Segment Tree でデータを持つと、上記の2操作がいずれも $\Theta (\text{log}W)$ （ $W$ : グリッドの横幅）で処理できます。

詳細は端折りますが、ざっくり概要だけ解説します。

Segment Tree は図のような完全二分木の構造をしており、葉の部分には平面走査のデータ、その他の節点には左右の子の和が保存されています。ここで、完全二分木の高さは $\Theta(\text{log} (W) )$ 、節点数は $\Theta (W)$ になることに注意してください。

このようにデータを持つと、1点更新は図のように $\Theta (\text{log} W)$ 個の節点の更新で済みます。

また、任意の区間和についても参照する節点数が $\Theta (\text{log} W)$ で抑えられることが示せます*6。

というわけで、Segment Tree を用いれば1点更新と区間和の取得の両方を高速に処理できます。

Segment Tree を用いた場合の計算量は、次の通りです：

平面走査については、長さの配列を用意したうえで更新が回発生するので、
- 黄色部分、緑部分の総和を求めるのは各 $i$ ごとに $\Theta (\color{red}{\text{log}W})$ で、まとめると $\Theta (\color{red}{Q\text{log}W})$
→ 判定問題 $P_i(d_i) ~ (1 \le i \le Q)$ をまとめて解く処理1回あたり $\Theta (W+N+\color{red}{Q\text{log}W})$
→ 全体 $\Theta ( (W+N+\color{red}{Q\text{log}W}) \text{log} D)$