【競プロ典型90問】007をPythonで解説

はじめに

最近競プロを始めたので，競プロ典型90問の★1~★3を解いて典型問題へのアプローチ方法を学んでいます．
公式のソースコードがC++で困っているPython勢の方も多いと思うので，自分の解答をシェアしたいと思います！
初心者の解答なので，もっといい方法があれば教えていただけると嬉しいです！

007 CP Classes（★3）

問題文

ABC 競技プログラミング塾には $N$ 個のクラスがあり、番号 $i$ $(1≤i≤N)$ のクラスはレーティング $A_i$ の生徒を対象としています。
今、$Q$ 人の生徒がこの塾に入塾しようとしています。番号 $j$ $(1≤j≤Q)$ の生徒のレーティングは $B_j$ です。各生徒は自分に合わないレベルのクラスに行くと不満になります。各生徒について、その不満度は次のように定義されます。

• 対象レーティング $a$ と自分のレーティング $b$ の差の絶対値、すなわち $∣a−b∣$

$j=1,2,3,…,Q$ それぞれについて、番号 $j$ の生徒の不満度として考えられる最小値を求めてください。ただし、$1$ 人も入らないクラス、複数人から成るクラスがあっても良いものとします。

制約

• $1≤N≤300000$
• $1≤Q≤300000$
• $0≤A_i≤10^9$
• $0≤B_j≤10^9$
• 入力はすべて整数

解法

$Q$ 人の生徒それぞれについて，$N$ 個のクラスの中で最も不満が少なくなるクラスを選べば良いです．
これを愚直に実装しようとすると，生徒 $1$ 人に対してその生徒のレーティング $b_j$ と $A_1 \sim A_N$ の $N$ 個の対象レーティングを比較してその最小値を求める操作を $Q$ 人の生徒に対して行う必要があります．この計算量は $O(NQ)$ で，$N, Q$ の制約からわかるようにTLE してしまいます．
そこで， 生徒 $1$ 人に対する操作を工夫することを考えます．
数値のリストの中から特定の値，もしくは最も近い値を探す方法として二分探索があります．ここでは二分探索について詳しくは説明しませんが，$1$ 回あたりの操作で探索範囲を半分にしていく探索方法で，リストの長さを $N$ とすると $O(\log N)$ の計算量で探索できます．
二分探索を行うためには，対象となるリストがソートされている必要があるので，まずは $A_1 \sim A_N$ を昇順ソートし，ソート後に $A_{1'} \sim A_{N'}$ になったものとします．そして，各生徒のレーティング $b_j$ について二分探索をして， $A_{1'} \sim A_{N'}$ の中にソートされた状態を保ったまま $b_j$ を挿入できる位置を探します．
ここで，もしその位置がリストの両端であれば， $b_j$ に最も近い値は隣の要素の値になります．一方で，リストの両端でない場合は，挿入できる位置の両隣の値のどちらかが最も $b_j$ の値に近くなるので，実際に計算してどちらのほうが不満度が小さくなるのかを比較する必要があります．
$A_1 \sim A_N$ のソートにかかる計算量は $O(N \log N)$ で，$Q$ 人の生徒に対する二分探索の計算量は $O(Q\log N)$ なので，全体で $O(( N+Q) \log N)$ の計算量で実装できます．

↓↓↓ 二分探索の説明はこちらがわかりやすいのでオススメです ↓↓↓

実装

Pythonで二分探索を行う際に使えるモジュールとして bisect があります．これは競プロをしているとよく使うモジュールだと思うので，使い方を覚えておくと良いと思います．公式リファレンスを確認してみてください．

answer.py

import sys
input = sys.stdin.readline
from bisect import bisect_left


N = int(input())
A = list(map(int, input().split()))
Q = int(input())
A.sort()
ans = []
for _ in range(Q):
    B = int(input())
    # ソートされた順序を保ったままBをAに挿入できる位置を二分探索
    i = bisect_left(A, B)
    if i == 0:
        # BがAのどの値よりも小さければ最も低いレベルのクラスに行く
        ans.append(abs(A[0] - B))
    elif i == N:
        # BがAのどの値よりも大きければ最も高いレベルのクラスに行く
        ans.append(abs(A[N-1] - B))
    else:
        # 挿入後の両端の不満度を比較して小さい方へ行く
        ans.append(min(abs(A[i-1] - B), abs(A[i] - B)))
# 出力
for i in range(Q):
    print(ans[i])

さいごに

今回使った二分探索は，競プロをしているといろいろなパターンの問題で出会うと思います．汎用性が高く，モジュールを使うことで簡単に実装できるので，この機会にしっかり身につけておくことをオススメします．
よく分からないことやコードの改善点などありましたら, 気軽にコメントしてください.

最後までお付き合いいただきありがとうございました！