ABC155 D - Pairs

問題原文

atcoder.jp

問題要旨

個の整数がある。 これらの整数のペアは、 通り考えられるが、それぞれのペアについて積を求め、小さい順に並び替えた時 番目にくる数は何か？

• 
• 

解法

atcoder.jp

こいつの上位互換問題。。。どう考えても令和ABC-Dの難易度じゃない。

難しくしているポイントは以下の２点。

• 負の数が混ざる
• 同じペアについて、順番を入れ替えただけのものをカウントできない

アプローチとしては、まず正の数の個数、0の個数、負の数の個数を出して、「 番目の値が正か０か負か」を先に求めてしまうと楽になる。

答えが負の場合

二分探索で、「積が 以下になるペアが 個以下あるか？」に答えていく。探したいのはギリギリ 個を達成するような で、これが 番目の値になる。
積が 以下になるペアを数えるのには二分探索を使えばいい。けどPythonだとこの制約で二分探索 in 二分探索をやるのは厳しいので、尺取り方でやる。詳しくは実装で。
答えが負になることが前提なので、負の要素と正の要素から１つずつとったもののみ考えればよくなっている。

答えが０の場合

これはそのまま０を出力して終わり。 平和な世界。

答えが正の場合

まず積が負になるようなペアと積が０になるようなペアの個数を から引いておく。 そのあとはこちらも二分探索で、「積が 以下になるペアが 個以下あるか？」に答えていく。探したいのはギリギリ 個を達成するような で、これが 番目の値になる。

積が 以下になるペアを数えるのには同じく二分探索を使えばいいが、Pythonだと厳しいので、尺取り方でやる。詳しくは実装で。
答えが正になることが前提なので、負の要素から２つ or 正の要素から２つといった選び方を考える。

実装

かなり頭が爆発する。
まずソートしておいて、負の要素の境界 と 正の要素の境界 を求めておく。

以下３種類の尺取り法について、簡単に説明する。

1. 答えが負の場合の尺取り 選び方は、負の要素から１つ、正の要素から１つ選ぶような選び方。
   例えば、 とする。 ある負の要素 と正の要素 に着目した時、 が 以下ならば、 を１つ右にずらした も必ず 以下になる。
   逆に考えれば が より大きければ、 を１つ右にずらした も必ず より大きくなるので 、そのようなものは調べる必要がない。

2. 答えが正の場合、正 × 正のカウント 正の数から２つ選ぶような選び方を考える。同じペアをダブルカウントしないように要注意。
   とする。 が 以下になった場合、 も必ず 以下になる。
   逆に考えれば、 が より大きいとき、 は も必ず より大きくなるので調べる必要がない。
   よって、 両端から真ん中に向けて詰めるように を動かしていく。 ダブルカウントしないためには、 が必ず成立するようにすればいい。

3. 答えが負の場合、負 × 負のカウント 2の逆パターンをやる。
   各 について、 がどこまでいけるかずらすようにするとやりやすい。好み。
   こちらも を条件におけばダブルカウントを防げる。

よって、負の数・正の数ともに小さい方から見ていけばいい。
ある が 以下となった場合、 残った右側のものがくっつきうる対象になる。

from bisect import bisect_left, bisect_right
N, K = map(int, input().split())
A = sorted(list(map(int, input().split())))
 
negative_end = bisect_left(A, 0)
positive_start = bisect_right(A, 0)
 
positive_cnt = N - positive_start
negative_cnt = negative_end
zero_cnt = N - (positive_cnt + negative_cnt)
 
negative_pairs_cnt = positive_cnt * negative_cnt
zero_pairs_cnt = zero_cnt * (positive_cnt + negative_cnt) + (zero_cnt * (zero_cnt - 1) // 2)


# 答えが負の場合
if K <= negative_pairs_cnt:
    # にぶたん: X以下になるペアがK個以上あるか？
    ok, ng = 0, -10 ** 18
    while abs(ok - ng) > 1:
        X = (ok + ng) // 2
        or_higher_cnt, r = 0, positive_start
        for l in range(negative_end):
            while (r < N) and (A[l] * A[r] > X):
                r += 1
            or_higher_cnt += positive_cnt - (r - positive_start)
 
        if or_higher_cnt >= K:
            ok = X
        else:
            ng = X
 
    print(ok)

# 答えが0の場合
elif K <= negative_pairs_cnt + zero_pairs_cnt:
    print(0)
 
# 答えが正の場合
else:
    # にぶたん: X以下になるペアがK個以上あるか？
    K -= (negative_pairs_cnt + zero_pairs_cnt)
    ok, ng = 10 ** 18, 0
    while abs(ok - ng) > 1:
        X = (ok + ng) // 2
        or_higher_cnt = 0
 
        r = N - 1
        for l in range(positive_start, N):
            while (r < N) and (A[l] * A[r] > X) and (l < r):
                r -= 1
            or_higher_cnt += max(0, r - l)
 
        l = 0
        for r in range(negative_end - 1, -1, -1):
            while (l < N) and (A[l] * A[r] > X) and (l < r):
                l += 1
 
            or_higher_cnt += max(0, r - l)
 
        if or_higher_cnt >= K:
            ok = X
        else:
            ng = X
 
    print(ok)

感想

絶対400違うやろお前ーーーーーーー！！！！！！！！！！！！！