[Python] 二分探索 ABC155D

ABC155D

小さい順に K 番目の値を求めよ
⇒　x 以下のものが K 個以上あるような、最小の x を求めよ

内容は下記の引用である。この投稿は私的メモである。
ABC155【D問題の復習】文字起こしをするくらいの勢いで解説動画を復習したら、二分探索の苦手感が減ってきた

AtCoder Beginner Contest 155(ABC155) A～E解説

サンプルコード

n,k=map(int,input().split())
arr=list(map(int,input().split()))
arr=sorted(arr)
nc=0
zc=0
pc=0

# 負数の数(nc)、0の数(zc)、正数の数(pc)を数える
for i in range(n):
 if arr[i]<0:
   nc+=1
 elif arr[i]==0:
   zc+=1
 else:
   pc+=1

# 積が負となるペアの個数を求める
ncc=nc*pc
# 積が0となるペアの個数を求める
zcc=zc*(nc+pc)+(zc*(zc-1))//2
# 積が正となるペアの個数を求める
pcc=(nc*(nc-1))//2+(pc*(pc-1))//2

# K番目の値が負の場合
if k <= ncc:
 arr1=[]
 arr2=[]
 # 負数、正数をそれぞれ昇順に並べる
 for i in range(n):
   if arr[i] < 0:
     arr1.append(-arr[i])
   elif arr[i] > 0:
     arr2.append(arr[i])
 arr1=arr1[::-1]
 c1=len(arr1)
 c2=len(arr2)
 l = 0         # left
 r = 10**18+1  # right 答えでもある
 # 負数のマイナスを取って計算しているため積の大小が逆転するので、大きなほうからK番目
 # 小さなほうから(積が負となるペアの個数-K)番目の値を求める
 k = ncc-k
 while r-l != 1: # rとlが1違い、すなわちrが答えとなるまで繰り返す
   mid=(l+r)//2 # 区間の中間値を更新
   cnt=0
   pos=c2-1
   # 区間[left, right)でのしゃくとり法
   # 各負数ごとに積が勝手に決めた値midよりも小さくなる正数の個数を求め足し合わせる
   for i in range(c1):
     while pos!=-1:
       if arr2[pos] > mid//arr1[i]:
         pos-=1
       else:
         break        
     cnt+=pos+1
   # 勝手に決めた値mid以下の個数が(積が負となるペアの個数-K)個以上であれば真の値はmid以下
   if cnt > k:
     r = mid  # rightをmidに更新してやり直し
   # そうでなければ真の値はmid以上
   else:
     l = mid  # leftをmidに更新してカウントやり直し
 # 負数のマイナスを取って計算したので最後にマイナスを付けて出力する
 print(-r)

# K番目の値が0の場合
elif k <= (ncc+zcc):
 print(0)

# K番目の値が正の場合
else:
 arr1=[]
 arr2=[]
 for i in range(n): #負数、正数をそれぞれ昇順に並べる
   if arr[i] < 0:
     arr1.append(-arr[i])
   elif arr[i] > 0:
     arr2.append(arr[i])
 arr1=arr1[::-1]
 c1=len(arr1)
 c2=len(arr2)
 l=0
 r=10**18+1
 k -= (ncc+zcc) #正数の中で小さい順にK番目の数を調べたいので負数、0の数を取り除く
 while r-l != 1:
   mid=(l+r)//2
   cnt1=0
   pos1=c1-1
   # 各負数ごとに積が勝手に決めた値midよりも小さくなる負数の個数を求め足し合わせる(しゃくとり法)
   for i in range(c1):
     tmp=0
     while pos1 != -1:
       if arr1[pos1] > mid//arr1[i]:
         pos1-=1
       else:
         break
     tmp+=pos1+1
     # 各要素について、自身を2回選んだ場合の積が勝手に決めた値mid以下なら条件を満たすペアの個数を1減らす
     if arr1[i]**2 < mid:
       tmp-=1
     cnt1+=tmp
   cnt1 = cnt1//2 #上記の処理では同じペアが2回数えられているので個数を半分にする
   cnt2=0
   pos2=c2-1
   # 各正数ごとに積が勝手に決めた値midよりも小さくなる正数の個数を求め足し合わせる(しゃくとり法)
   for i in range(c2):
     tmp=0
     while pos2!=-1:
       if arr2[pos2] > mid//arr2[i]:
         pos2-=1
       else:
         break
     tmp+=pos2+1
     if arr2[i]**2 < mid:
       tmp-=1 #各要素について、自身を2回選んだ場合の積が勝手に決めた値mid以下なら条件を満たすペアの個数を1減らす
     cnt2+=tmp
   cnt2=cnt2//2 #上記の処理では同じペアが2回数えられているので個数を半分にする
   cnt=cnt1+cnt2
   if cnt >= k: #勝手に決めた値mid以下の個数がK個以上であれば真の値はmid以下
     r = mid
   else: #そうでなければ真の値はmid以上
     l = mid
 print(r)