ABC155-D_Pairsをpythonで解いた（二分探索）

問題

ABC155-D_Pairs

ポイント

小さい順に並べ替えたとき、K番目にくる数...
-> ある条件を満たすものの最小値（最大値）を求める
-> 二分探索
- AtcorderABC-Dより先へ行くために代表的なアルゴリズムを勉強した-その1-（DP、累積和、UnionFind木、bit関係、二分探索）

def binary_search(sorted_list, n):
    '''
    配列に対してある数nを配列のソート順を崩さないように挿入するためのindexを返す関数。
    同一の数が配列に複数含まれていた場合は最も小さいindexを返す。
    '''
    min = 0
    max = len(sorted_list) - 1
    while min <= max:
        mid = min + (max - min) // 2
        if sorted_list[mid] == n:
            if sorted_list[mid - 1] == n:
                max = mid
            else:
                return mid
        if sorted_list[mid] < n:
            min = mid + 1
        else:
            max = mid - 1
    return min

example_list = [1,3,4,5,6,7,9,10,13,16,18,22,24,25,26] 
target = 8

example_list = sorted(example_list)
binary_search(example_list, target)
# -> 6

解説

積が負のペア，0 のペア，正のペアの各個数は簡単に求まるので，答えが負・0・正のどれになるか
はこれでわかります．
答えが負になる場合は，負の数と正の数を 1 つずつ選んで x 以上になるペアがいくつあるかを数
えることが尺取り法によって可能なので，二分探索を用いて答えが求まります．
答えが正になる場合も全く同様ですが，同じ要素を 2 回選ぶこと，それを差し引くと各ペアをちょ
うど 2 回数えることを考慮しましょう．

解法

n ,k = map(int,input().split())
a_list = [int(i) for i in input().split()]

# プラスとマイナスに分ける
plus_list = []
minus_list = []
zero = 0
for i in range(n):
    if a_list[i] > 0:
        plus_list.append(a_list[i])
    elif a_list[i] < 0:
        minus_list.append(-a_list[i])
    else:
        zero += 1

plus_list.sort()
minus_list.sort()

# 二分探索
ng = -(10 ** 18)-1
ok = 10 ** 18+1

while ok - ng > 1:
    mid = (ok + ng) // 2
    s = 0
    if mid < 0:
        j = 0
        for i in range(len(minus_list) - 1, -1, -1):
            while j < len(plus_list) and minus_list[i] * plus_list[j] < -mid:
                j += 1
            s += len(plus_list) - j

    else :
        j = len(plus_list) - 1
        for i in range(len(plus_list)):
            while j >= 0 and plus_list[i] * plus_list[j] > mid:
                j -= 1
            s += max(0, j - i)

        j = len(minus_list) - 1
        for i in range(len(minus_list)) :
            while j >= 0 and minus_list[i] * minus_list[j] > mid :
                j -= 1
            s += max(0, j - i)

        s += (len(plus_list) + len(minus_list)) * zero + len(plus_list) * len(minus_list) + zero * (zero - 1) // 2

    if s < k :
        ng = mid
    else :
        ok = mid

print(ok)