二分探索とABC205 D

更新

二分探索に関する便利なライブラリ bisect を見つけたので更新しました.

概要

ABC205のD問題解き方はわかったのに, 普段とは違う二分探索の実装に脳みそが追いつかずハチャメチャになってコンテスト中にACにならなかった....

ということで, ABC205の解説とそれに関する二分探索についてまとめる
https://atcoder.jp/contests/abc205

D - Kth Excluded

配列 $A = [A_1, A_2, ... ,A_n]$とInt $k (1\leq k \leq 10^{18})$が与えられる.
このときに, $A$に含まれない正の整数で, $k$番目の数を答える問題.

公式解説に習って, $A$に含まれない正の整数を"良い数字"と呼ぶことにする.

例

$A = [3, 5, 6, 7]$と$k = 5$が与えられたとする.
$A$に含まれない正の整数は,
$1, 2, 4, 8, 9, 10, ....$ なので, 答えである$5$番目の数は$9$

解説

$k$の値が非常に大きくなるため, $A$に含まれない配列を予め全て計算して答えを求めるような方法は明らかに時間が足りない.

そこで, $k$番目の良い数字が配列$A$のどことどこの間に含まれるかにまずは注目する.

この配列$A$のどことどこの間に含まれるか & その間において何番目の数字か というのがわかれば, 配列$A$への1回のアクセスと何回かの演算で答えが求まりそう.

どことどこの間に含まれるか

これを求めるために, 配列$A$の各要素より前に良い数字がいくつ存在するかという配列を事前に求めておく. この配列をgood_numと名付けておきます.

これを求めるのは簡単で, 各$i(1\leq i \leq n)$について, $(A_{i} - A_{i-1} - 1)$を計算すれば良いです.

気をつけるのは, $i=1$のときには$A_i - 1$を計算することと, 各要素より前に（総じて）いくつ存在するかの和を計算することです.

good_num = [0 for i in range(len(A))]
good_num[0] = A[0] - 1
for i in range(1, len(A)):
  good_num[i] = good_num[i-1] + (A[i] - A[i-1] - 1)

この配列が求まった後に, どことどこに求める良い数字があるのか の条件は, $k$以上となる最小のIndexを$s$として, $A_s$と$A_{s-1}$の間にあります.

その間において何番目か

この計算は簡単で, 今$A_i$より前に$s(\geq k)$個の良い値が存在したとする.

すると実際の良い値は, $A_i - 1 - (s - k)$となる.
わからない人は数直線上で確認してみて.

例外の処理

$A_n$よりも右側に求める良い値が存在する場合 と $A_1$よりも左側に求める良い値が存在する場合は少し怖いので例外的な処理をしておきます.

• 例外1 : $A_n$よりも右側にある場合
  このとき, $A_n$よりも前に存在する良い値の個数を$s$とすると, 求める良い値は$A_n + (k - s)$となります. 数直線で確認してみて.

• 例外2 : $A_1$よりも左側にある場合
  このときは単純に$k$が答えになります.

直線探索による実装

解く方針は固まったので, 配列good_numにおいて$k$以上になる最小のIndexを見つける方法に直線探索を使ってといてみます.

def linear_search(data, k):
  for i in range(len(data)):
    if data[i] >= k:
      return i

for _ in range(Q):
  k = int(input())
  # 例外1の処理
  if good_num[-1] < k:
    ans = A[-1] + (k - good_num[-1])
  else:
    s = linear_search(good_num, k)
    ans = (A[s] - 1) - (good_num[s] - k)
  print(ans)

このコードは正しい答えを返してくれますが, 計算量が$O(Q * N)$になってしまいます.
もちろんTLEです. Pypyでも.

二分探索の話

上の直線探索の部分を二分探索に改良できれば, 計算量が$O(Q * \log N)$となってなんだか解けそうです.

求める値と一致する配列のIndexを返す二分探索

一致する値がない場合は-1を返します. よくある二分探索

def binary_search(data, value):
  left = 0
  right = len(data) - 1
  while right - left > 1:
    mid = (left + right) // 2
    if data[mid] == value:
      return mid
    elif data[mid] > value:
      right = mid
    else:
      left = mid
  return -1

条件を満たす最小のIndexを返す二分探索

今この記事を書いてて図とかはったらわかりやすいと思ったのですが, 図を作るのがめんどくさいのでこの記事とかを参考にしてください.
https://qiita.com/drken/items/97e37dd6143e33a64c8c

結果だけ書くと

def binary_search2(data, value):
  left = 0
  right = len(data) - 1
  while right - left > 1:
    mid = (left + right) // 2
    if data[mid] >= value:
      right = mid
    else:
      left = mid
  return right

rightには常に条件を満たす様な配列のIndexを, leftには常に条件を満たさないような配列のIndexを持っておきます.
それで最終的にrightを出力します.

ACした解答

# 二分探索
def binary_search(data, k):
  left = 0
  right = len(data) - 1
  while right - left > 1:
    mid = (left + right) // 2
    if data[mid] >= k:
      right = mid
    else:
      left = mid
  return right

# 入力の受け取りと前処理
N, Q = map(int, input().split())
A = list(map(int, input().split()))
good_num = [0 for i in range(len(A))]
good_num[0] = A[0] - 1
for i in range(1, len(A)):
  good_num[i] = good_num[i-1] + (A[i] - A[i-1] - 1)

# 実際の処理
for _ in range(Q):
  k = int(input())
  # 例外2の処理
  if good_num[0] >= k:
    ans = k
  # 例外1の処理
  elif good_num[-1] < k:
    ans = A[-1] + (k - good_num[-1])
  # それ以外
  else:
    s = binary_search(good_num, k)
    ans = (A[s] - 1) - (good_num[s] - k)
  print(ans)

追記：二分探索 bisect

Pythonの標準ライブラリに便利なモノを見つけた. 実際に二分探索を目的としたライブラリではないらしいのだが, 二分探索を支援するライブラリらしい（中身は二分探索で実装されていることを確認した）.

ソートされたリストの中で, 指定の値をどのインデックスに挿入すれば, ソートされた状態が崩れないで挿入できるか, というのに使うのがbisect.bisect() である.
実際に試してみる.

import bisect

A = [1,2,3,3,3,3,4,5,6]
# 3をAのどこに挿入したら良いか, 一番左のindexを返す
index = bisect.bisect_left(A, 3)
print(index)
>>> 2

# 一番右のindexを返す
index = bisect.bisect_right(A, 3)
print(index)
>>> 6

これを使えば, $k$以上となる最小のindexはbisect.bisect_left()を使えば見つかる.
以下このbisectを使った実装

import bisect
# 入力の受け取りと前処理
N, Q = map(int, input().split())
A = list(map(int, input().split()))
good_num = [0 for i in range(len(A))]
good_num[0] = A[0] - 1
for i in range(1, len(A)):
  good_num[i] = good_num[i-1] + (A[i] - A[i-1] - 1)

# 実際の処理
for _ in range(Q):
  k = int(input())
  # 例外2の処理
  if good_num[0] >= k:
    ans = k
  # 例外1の処理
  elif good_num[-1] < k:
    ans = A[-1] + (k - good_num[-1])
  # それ以外
  else:
    # 変更点, 探索にbisectを使う
    s = bisect.bisect_left(good_num, k)
    ans = (A[s] - 1) - (good_num[s] - k)
  print(ans)

これ便利ですね, 配列に指定した値がいくつかるか とか 普通の二分探索もこれを使えば良さそうです.