PythonでAtCoder（ABC364D問題）解いてみた

はじめに

問題はこちら
初心者（灰色〜茶色）向けです。
D問題の公式解説を咀嚼してPythonで実装したものについて説明します。

D - K-th Nearest

考え方

大枠の方針

二分探索を用います。$k$番目に大きい/小さい値を求める際に有効な手段です。
今回、次の単調増加な関数$f_j$を用いると、各$x$について、$f_j(x)$たちはソート済の列と見なすことができ、その中で$k_j$番目に小さいものを探せばよいということになります。¹

詳細

$x\geq0$ について、

f_j(x) := B_j\spaceを中心とした長さxの区間に含まれる点\space A_1, A_2, \ldots, A_N \spaceの数

と定めます。²
このとき、 $f_j(x)$ は$x$に関して（広義）単調増加です。³
$x$の値を変化させていって、$f_j(x)\geq k_j$となる$x$の最小値を求めれば良いです。⁴
なお、$f_j(x)$の値が変わるのは、区間が各$A_i$の点を跨ぐとき、つまり、$x=|b_j-a_i|$のときです。

また、$f_j$は単調なので、二分探索により、出力する$x$の値を広い範囲から狭めて求めることが可能です。
問題の制約上、今回の最小値は$0$、最大値は$2*10^8$ですから、二分探索による計算量は$ \log(2 * 10^8)$程度です。

あとは、$f_j(x)$を実際に構成できればよいということになりますが、$a_n$をソートしておくと、二分探索を用いて下記のように構成できます。

f_j(x) の構成の仕方

$f_j(x)$の実装

def f(x) :
    indexL = bisect.bisect_left(a_n,b-x)
    indexR = bisect.bisect_right(a_n,b+x)
    return indexR-indexL

$f_j(x)$　は下記のように言い換えられます。

\begin{align}
f_j(x) &= 区間[b_j-x,b_j+x]に含まれる A_n たちの数 \\
       &= (b_j+x含めそれより左側のA_nたちの数)-(b_j-xより左側のA_nたちの数)
\end{align}

ここで、

\begin{align}
(b_j+x含めそれより左側のA_nたちの数) &= b_j+xを a_nで二分探索した右側の挿入位置　\\
(b_j-xより左側のA_nたちの数) &= b_j-xを a_nで二分探索した左側の挿入位置
\end{align}

ですから、bisectモジュールで計算可能です。計算量は $ \log(n)$です。

解答例

私の解答

import bisect
def f(x) :
    indexL = bisect.bisect_left(a_n,b-x)
    indexR = bisect.bisect_right(a_n,b+x)
    return indexR-indexL

n,q = map(int,input().split())
a_n = sorted(list(map(int,input().split())))
for _ in range(q) :
    b,k = map(int,input().split())
    boundL = -1
    boundR = 2*(10**8)
    cnt = 0
    mid = 0
    while boundR-boundL >1:
        mid = (boundL+boundR)//2
        if f(mid) >= k :
            boundR = mid
        else :
            boundL = mid
        cnt+=1
    print(boundR)

感想

単調な関数を用意すれば、k番目に大きい値は2分探索を用いて算出できるということを学びました。
配列でなくてもよい、アルゴリズムを自分で書くケースがあるのだと知ることができました。

しばらく投稿をサボっていましたが、復活します。できなかった問題、面白いと思ったものだけピックアップして記載します。
文章を書くのは難しいですね。

参考

二分探索
https://qiita.com/Pro_ktmr/items/8946723fe08ba29a977c

1. 「$k_j$番目に小さいもの」は厳密ではありません。 $f_j(x) \geq k_j$を満たす$x$の最小値です。 ↩

2. $B_j$を中心とした長さ$x$の区間 とは、 $[b_j-x,b_j+x]$ のことです。 ↩

3. 単調増加の理由は、$x$が小さくなるほど区間が狭くなり、区間に含まれる点の数が小さくなるためです。また、$f_j(x)$は$x=0$の時最小になり、$x$を十分に大きくすれば、最大値の$n$になります。 ↩

4. $f_j(x)= k_j$　でなく $f_j(x)\geq k_j$ なのは、$A_i$たちには重複があり得るので、ちょうど$k_j$になる$x$が存在するとは限らないからです。 ↩