AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC155 D - Pairs

問題原文

atcoder.jp

問題要旨

 N 個の整数がある。 これらの整数のペアは、  \frac{N(N - 1)}{2} 通り考えられるが、それぞれのペアについて積を求め、小さい順に並び替えた時  K 番目にくる数は何か?

  •   1 \leq N \leq 2×10^{5}
  •  -10^{9} \leq a_i \leq 10^{9}

解法

atcoder.jp

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

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

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

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

答えが負の場合

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

答えが0の場合

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

答えが正の場合

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

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

実装

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

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

  1. 答えが負の場合の尺取り 選び方は、負の要素から1つ、正の要素から1つ選ぶような選び方。
    例えば、  A = [-5, -4, -3, 3, 4, 5\rbrack とする。 ある負の要素  negative_l と正の要素  positive_r に着目した時、  negative_l × positive_r  X 以下ならば、  r を1つ右にずらした  negative_l × positive_{r + 1} も必ず  X 以下になる。
    逆に考えれば  negative_l × positive_r X より大きければ、  l を1つ右にずらした   [tex: negative_{l + 1} × positive_r も必ず  X より大きくなるので 、そのようなものは調べる必要がない。

  2. 答えが正の場合、正 × 正のカウント 正の数から2つ選ぶような選び方を考える。同じペアをダブルカウントしないように要注意。
     A = [1, 2, 3, 4, 5\rbrack とする。 a_l × a_r X 以下になった場合、  a_l × a_{r - 1} も必ず  X 以下になる。
    逆に考えれば、 a_l × a_r X より大きいとき、 a_{l + 1} × a_r X も必ず  X より大きくなるので調べる必要がない。
    よって、 両端から真ん中に向けて詰めるように  l, r を動かしていく。 ダブルカウントしないためには、  l < r が必ず成立するようにすればいい。

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

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

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違うやろお前ーーーーーーー!!!!!!!!!!!!!