Problem p02774 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define N 100005
typedef long long ll;
using namespace std;

int a[N], b[N];
vector<int> v1, v2;

ll k;

bool check1(ll mid) {
  ll num = 0;
  for (int i = 0; i < v1.size(); i++) {
    ll x = mid / v1[i];
    if (mid % v1[i])
      x++;
    int p = lower_bound(v1.begin(), v1.end(), x) - v1.begin();
    p = max(i + 1, p);
    num += v1.size() - p;
  }
  // cout<<num<<endl;
  for (int i = 0; i < v2.size(); i++) {
    ll x = mid / v2[i];
    if (mid % v2[i])
      x++;
    int p = lower_bound(v2.begin(), v2.end(), x) - v2.begin();
    p = max(i + 1, p);
    num += v2.size() - p;
  }
  // cout<<num<<endl;
  return (num >= k);
}

bool check2(ll mid) {
  ll num = 0;
  for (int i = 0; i < v1.size(); i++) {
    ll x = mid / v1[i];
    if (mid % v1[i])
      x++;
    int p = lower_bound(v2.begin(), v2.end(), x) - v2.begin();
    num += v2.size() - p;
  }
  return (num >= k);
}

int main() {
  int n, m;
  scanf("%d%lld", &n, &k);
  k = 1ll * (n - 1) * n / 2 - k + 1;
  // cout<<k<<endl;
  for (int i = 1; i <= n; i++)
    scanf("%d", &a[i]);
  int a0 = 0;
  for (int i = 1; i <= n; i++) {
    if (a[i] > 0) {
      v1.push_back(a[i]);
    } else if (a[i] < 0) {
      v2.push_back(-a[i]);
    } else {
      a0++;
    }
  }

  ll sum0 = 1ll * a0 * (n - a0) + 1ll * a0 * (a0 - 1) / 2;
  ll sum1 = 1ll * (int)v1.size() * ((int)v1.size() - 1) / 2 +
            1ll * (int)v2.size() * ((int)v2.size() - 1) / 2;
  ll sum2 = 1ll * (int)v2.size() * (int)v1.size();
  // cout<<sum1<<" "<<sum0<<" "<<sum2<<endl;
  sort(v1.begin(), v1.end());
  sort(v2.begin(), v2.end());

  if (k <= sum1) {
    // cout<<check1(10)<<endl;
    ll l = 1, r = 1e18, ans = 1;
    while (l <= r) {
      ll mid = (l + r) / 2;
      if (check1(mid))
        l = (ans = mid) + 1;
      else
        r = mid - 1;
    }
    printf("%lld\n", ans);
  } else if (k > sum1 && k <= sum0 + sum1) {
    puts("0");
  } else {
    k = sum0 + sum1 + sum2 - k + 1;
    ll l = -1e18, r = -1, ans = -1;
    while (l <= r) {
      ll mid = (l + r) / 2;
      if (check2(-mid))
        r = (ans = mid) - 1;
      else
        l = mid + 1;
    }
    printf("%lld\n", ans);
  }
}
/*
3 3 3
2 3 4
4 5 6

3 3 7
-2 -3 -4
4 5 6
*/


Fixed submission:
#include <bits/stdc++.h>
#define N 200005
typedef long long ll;
using namespace std;

int a[N], b[N];
vector<int> v1, v2;

ll k;

bool check1(ll mid) {
  ll num = 0;
  for (int i = 0; i < v1.size(); i++) {
    ll x = mid / v1[i];
    if (mid % v1[i])
      x++;
    int p = lower_bound(v1.begin(), v1.end(), x) - v1.begin();
    p = max(i + 1, p);
    num += v1.size() - p;
  }
  // cout<<num<<endl;
  for (int i = 0; i < v2.size(); i++) {
    ll x = mid / v2[i];
    if (mid % v2[i])
      x++;
    int p = lower_bound(v2.begin(), v2.end(), x) - v2.begin();
    p = max(i + 1, p);
    num += v2.size() - p;
  }
  // cout<<num<<endl;
  return (num >= k);
}

bool check2(ll mid) {
  ll num = 0;
  for (int i = 0; i < v1.size(); i++) {
    ll x = mid / v1[i];
    if (mid % v1[i])
      x++;
    int p = lower_bound(v2.begin(), v2.end(), x) - v2.begin();
    num += v2.size() - p;
  }
  return (num >= k);
}

int main() {
  int n, m;
  scanf("%d%lld", &n, &k);
  k = 1ll * (n - 1) * n / 2 - k + 1;
  // cout<<k<<endl;
  for (int i = 1; i <= n; i++)
    scanf("%d", &a[i]);
  int a0 = 0;
  for (int i = 1; i <= n; i++) {
    if (a[i] > 0) {
      v1.push_back(a[i]);
    } else if (a[i] < 0) {
      v2.push_back(-a[i]);
    } else {
      a0++;
    }
  }

  ll sum0 = 1ll * a0 * (n - a0) + 1ll * a0 * (a0 - 1) / 2;
  ll sum1 = 1ll * (int)v1.size() * ((int)v1.size() - 1) / 2 +
            1ll * (int)v2.size() * ((int)v2.size() - 1) / 2;
  ll sum2 = 1ll * (int)v2.size() * (int)v1.size();
  // cout<<sum1<<" "<<sum0<<" "<<sum2<<endl;
  sort(v1.begin(), v1.end());
  sort(v2.begin(), v2.end());

  if (k <= sum1) {
    // cout<<check1(10)<<endl;
    ll l = 1, r = 1e18, ans = 1;
    while (l <= r) {
      ll mid = (l + r) / 2;
      if (check1(mid))
        l = (ans = mid) + 1;
      else
        r = mid - 1;
    }
    printf("%lld\n", ans);
  } else if (k > sum1 && k <= sum0 + sum1) {
    puts("0");
  } else {
    k = sum0 + sum1 + sum2 - k + 1;
    ll l = -1e18, r = -1, ans = -1;
    while (l <= r) {
      ll mid = (l + r) / 2;
      if (check2(-mid))
        r = (ans = mid) - 1;
      else
        l = mid + 1;
    }
    printf("%lld\n", ans);
  }
}
/*
3 3 3
2 3 4
4 5 6

3 3 7
-2 -3 -4
4 5 6
*/