Problem p02774 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <queue>
#include <string>
#include <vector>
using namespace std;
long long tot1, tot2, a[200010], b[200010], c[200010], kk;
long long check1(long long ys) {
  long long ans = 0;
  for (long long i = 1; i <= tot1; i++) {
    long long k = 1, s = tot2 + 1;
    while (k) {
      if (s > 0 && b[i] * c[s - k] < ys)
        s -= k, k = k * 2;
      else
        k = k / 2;
    }
    ans += tot2 - s + 1;
  }
  return ans;
}
long long calc1() {
  long long l = -2e18, r = -1, aa = 0;
  while (l <= r) {
    long long mid = (l + r) / 2, x = check1(mid), y = check1(mid + 1);
    if (x < kk && kk <= y)
      return mid;
    else if (x < kk)
      l = mid + 1;
    else
      r = mid - 1;
  }
}
long long check2(long long ys) {
  long long ans = 0;
  for (long long i = 1; i <= tot1; i++) {
    long long k = 1, s = i;
    while (k) {
      if (s + k <= tot1 && b[i] * b[s + k] < ys)
        s += k, k = k * 2;
      else
        k = k / 2;
    }
    ans += s - i;
  }
  for (long long i = 1; i <= tot2; i++) {
    long long k = 1, s = i;
    while (k) {
      if (s + k <= tot2 && c[i] * c[s + k] < ys)
        s += k, k = k * 2;
      else
        k = k / 2;
    }
    ans += s - i;
  }
  return ans;
}
long long calc2() {
  long long l = 1, r = 2e18, aa = 0;
  while (l <= r) {
    long long mid = (l + r) / 2, x = check2(mid), y = check2(mid + 1);
    if (x < kk && kk <= y)
      return mid;
    else if (x < kk)
      l = mid + 1;
    else
      r = mid - 1;
  }
}
int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  long long n, ans1, ans2, ans3, s1 = 0, s2 = 0, s3 = 0;
  cin >> n >> kk;
  for (long long i = 1; i <= n; i++)
    cin >> a[i];
  for (long long i = 1; i <= n; i++)
    if (a[i] == 0)
      s1++;
    else if (a[i] < 0)
      s2++, b[++tot1] = a[i];
    else
      s3++, c[++tot2] = a[i];
  ans1 = s2 * s3;
  ans2 = s1 * (n - s1) + s1 * (s1 - 1) / 2;
  ans3 = s2 * (s2 - 1) / 2 + s3 * (s3 - 1) / 2;
  sort(b + 1, b + tot1 + 1);
  sort(c + 1, c + tot2 + 1);
  reverse(b + 1, b + tot1 + 1);
  if (ans1 < kk && kk <= ans1 + ans2) {
    cout << 0 << "\n";
    return 0;
  } else if (kk <= ans1) {
    long long ans = calc1();
    cout << ans << "\n";
    return 0;
  } else if (kk > ans1 + ans2) {
    kk = kk - ans1 - ans2;
    long long ans = calc2();
    cout << ans << "\n";
    return 0;
  }
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <queue>
#include <string>
#include <vector>
using namespace std;
long long tot1, tot2, a[200010], b[200010], c[200010], kk;
long long check1(long long ys) {
  long long ans = 0;
  for (long long i = 1; i <= tot1; i++) {
    long long k = 1, s = tot2 + 1;
    while (k) {
      if (s - k > 0 && b[i] * c[s - k] < ys)
        s -= k, k = k * 2;
      else
        k = k / 2;
    }
    ans += tot2 - s + 1;
  }
  return ans;
}
long long calc1() {
  long long l = -2e18, r = -1, aa = 0;
  while (l <= r) {
    long long mid = (l + r) / 2, x = check1(mid), y = check1(mid + 1);
    if (x < kk && kk <= y)
      return mid;
    else if (x < kk)
      l = mid + 1;
    else
      r = mid - 1;
  }
}
long long check2(long long ys) {
  long long ans = 0;
  for (long long i = 1; i <= tot1; i++) {
    long long k = 1, s = i;
    while (k) {
      if (s + k <= tot1 && b[i] * b[s + k] < ys)
        s += k, k = k * 2;
      else
        k = k / 2;
    }
    ans += s - i;
  }
  for (long long i = 1; i <= tot2; i++) {
    long long k = 1, s = i;
    while (k) {
      if (s + k <= tot2 && c[i] * c[s + k] < ys)
        s += k, k = k * 2;
      else
        k = k / 2;
    }
    ans += s - i;
  }
  return ans;
}
long long calc2() {
  long long l = 1, r = 2e18, aa = 0;
  while (l <= r) {
    long long mid = (l + r) / 2, x = check2(mid), y = check2(mid + 1);
    if (x < kk && kk <= y)
      return mid;
    else if (x < kk)
      l = mid + 1;
    else
      r = mid - 1;
  }
}
int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  long long n, ans1, ans2, ans3, s1 = 0, s2 = 0, s3 = 0;
  cin >> n >> kk;
  for (long long i = 1; i <= n; i++)
    cin >> a[i];
  for (long long i = 1; i <= n; i++)
    if (a[i] == 0)
      s1++;
    else if (a[i] < 0)
      s2++, b[++tot1] = a[i];
    else
      s3++, c[++tot2] = a[i];
  ans1 = s2 * s3;
  ans2 = s1 * (n - s1) + s1 * (s1 - 1) / 2;
  ans3 = s2 * (s2 - 1) / 2 + s3 * (s3 - 1) / 2;
  sort(b + 1, b + tot1 + 1);
  sort(c + 1, c + tot2 + 1);
  reverse(b + 1, b + tot1 + 1);
  if (ans1 < kk && kk <= ans1 + ans2) {
    cout << 0 << "\n";
    return 0;
  } else if (kk <= ans1) {
    long long ans = calc1();
    cout << ans << "\n";
    return 0;
  } else if (kk > ans1 + ans2) {
    kk = kk - ans1 - ans2;
    long long ans = calc2();
    cout << ans << "\n";
    return 0;
  }
  return 0;
}