atcoder 155 d問題が解けない

Question

atcoder155 d問題
https://atcoder.jp/contests/abc155/tasks/abc155_d
が解けません。解説動画を見て、そのコードは理解したのですが、なぜ自分のコードがだめなのかわかりません。

どこがだめなのかわからないので、教えてもらえると嬉しいです。

#include<iostream>
#include<cmath>
#include<algorithm>
#include<vector>
#define rep(i, n) for (int i=0; i<(n); i++)
typedef long long ll;
using namespace std;


int main(){
    int N, K; cin >> N >> K;
    vector<ll> A(N, 0);
    int minus = 0, plus = 0;
    rep(i, N) cin >> A[i];
    sort(A.begin(), A.end());
    rep(i, N){
        if(A[i] > 0) plus++;
        else if(A[i] < 0) minus++;
    }
    ll a = max(abs(A[0]), abs(A[N-1]));
    ll l = -a*a, r = a*a;
    ll tot=0;
    ll mid;
    while(l+1 < r){
        ll limit_minus = 0;
        ll limit_plus = N-1;
        mid = (l+r)/2;
        rep(i, N){
            if(A[i]<0){
                while(A[i]*A[limit_minus] > mid && limit_minus < N) limit_minus++;
                tot += N-limit_minus;
            }else{
                while(A[i]*A[limit_plus] > mid && limit_plus >= 0 ) limit_plus--;
                tot += limit_plus+1;
            }
        }
        rep(i, N){
            if(A[i]*A[i] <= mid) tot--;
        }
        tot /= 2;
        if(tot < K) l = mid;
        else r = mid;
        tot = 0;
    }
    cout << r << endl;
}

Comments

[kazuma-s]

plus と minus を求めていますが、その値を使っていません。説明の追加をお願いします。

Answer 1

            if(A[i]<0){
                while(A[i]*A[limit_minus] > mid && limit_minus < N) limit_minus++;
                tot += N-limit_minus;
            }else{
                while(A[i]*A[limit_plus] > mid && limit_plus >= 0 ) limit_plus--;
                tot += limit_plus+1;
            }

二分探索中のこの部分について、A = {-2, -1, 1, 2}, mid = 2 として考えてみると、
i = 0 の時、limit_minusは1となりtotに3が足されます。
i = 1 の時、limit_minusは1のままでtotに3が足されます。
しかし、i = 1の時にはtotに4が足されるべきです(Aの要素すべてに-1をかけると{2, 1, -1, -2}で、mid以下である要素の個数は4)。

これと同じことがmidがマイナスの場合のlimit_plusの動きにも当てはまります。mid = -2とすると、
i = 2 の時、limit_plusは0となりtotに1が足されます。
i = 3 の時、limit_plusは0のままでtotに1が足されます。
しかしi = 3の時にはtotに2が足されるべきです。

---

一番ネックなのはここだと思いますが、ほかにも多少問題があります。
一つは上のwhileの条件式で、インデックスの境界チェックはインデックスでアクセスする前にすべきです。境界チェックでランタイムエラーになる可能性があります。

while(A[i]*A[limit_minus] > mid && limit_minus < N) limit_minus++;
↓
while(limit_minus < N && A[i]*A[limit_minus] > mid) limit_minus++;

もう一つはKの値が32bitには収まらない可能性があることです。