Problem p02774 (C++)

Buggy submission (status: Time Limit Exceeded):
// #define NDEBUG
#include "bits/stdc++.h"
#include <array>
#include <iostream>
#include <random>
#include <string>
#include <unordered_map>
#include <unordered_set>
#ifdef _MSC_VER
#include <ppl.h>
// #include <boost/multiprecision/cpp_dec_float.hpp>
// #include <boost/multiprecision/cpp_int.hpp>
// #include <boost/rational.hpp>
// #else
// #include <omp.h>
#endif

using namespace std;

#define DUMPOUT cerr
#define dump(...)                                                              \
  DUMPOUT << "  ";                                                             \
  DUMPOUT << #__VA_ARGS__ << " :[" << __LINE__ << ":" << __FUNCTION__ << "]"   \
          << endl;                                                             \
  DUMPOUT << "    ";                                                           \
  dump_func(__VA_ARGS__)

using uint = unsigned;
using ll = long long;
using ull = unsigned long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using pdd = pair<double, double>;
using pss = pair<string, string>;
template <typename _KTy, typename _Ty>
ostream &operator<<(ostream &o, const pair<_KTy, _Ty> &m) {
  o << "{" << m.first << ", " << m.second << "}";
  return o;
}
template <typename _KTy, typename _Ty>
ostream &operator<<(ostream &o, const map<_KTy, _Ty> &m) {
  if (m.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *m.begin();
  for (auto itr = ++m.begin(); itr != m.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _KTy, typename _Ty>
ostream &operator<<(ostream &o, const unordered_map<_KTy, _Ty> &m) {
  if (m.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *m.begin();
  for (auto itr = ++m.begin(); itr != m.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const vector<_Ty> &v) {
  if (v.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << v.front();
  for (auto itr = ++v.begin(); itr != v.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const deque<_Ty> &v) {
  if (v.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << v.front();
  for (auto itr = ++v.begin(); itr != v.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const set<_Ty> &s) {
  if (s.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *(s.begin());
  for (auto itr = ++s.begin(); itr != s.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty>
ostream &operator<<(ostream &o, const unordered_set<_Ty> &s) {
  if (s.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *(s.begin());
  for (auto itr = ++s.begin(); itr != s.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const stack<_Ty> &s) {
  if (s.empty()) {
    o << "{ }";
    return o;
  }
  stack<_Ty> t(s);
  o << "{" << t.top();
  t.pop();
  while (!t.empty()) {
    o << ", " << t.top();
    t.pop();
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const list<_Ty> &l) {
  if (l.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << l.front();
  for (auto itr = ++l.begin(); itr != l.end(); ++itr) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _KTy, typename _Ty>
istream &operator>>(istream &is, pair<_KTy, _Ty> &m) {
  is >> m.first >> m.second;
  return is;
}
template <typename _Ty> istream &operator>>(istream &is, vector<_Ty> &v) {
  for (size_t t = 0; t < v.size(); t++)
    is >> v[t];
  return is;
}
template <typename _Ty> istream &operator>>(istream &is, deque<_Ty> &v) {
  for (size_t t = 0; t < v.size(); t++)
    is >> v[t];
  return is;
}
namespace aux { // print tuple
template <typename Ty, unsigned N, unsigned L> struct tp {
  static void print(ostream &os, const Ty &v) {
    os << get<N>(v) << ", ";
    tp<Ty, N + 1, L>::print(os, v);
  }
};
template <typename Ty, unsigned N> struct tp<Ty, N, N> {
  static void print(ostream &os, const Ty &v) { os << get<N>(v); }
};
} // namespace aux

template <typename... Tys>
ostream &operator<<(ostream &os, const tuple<Tys...> &t) {
  os << "{";
  aux::tp<tuple<Tys...>, 0, sizeof...(Tys) - 1>::print(os, t);
  os << "}";
  return os;
}

template <typename A, size_t N, typename T>
inline void Fill(A (&array)[N], const T &val) {
  fill((T *)array, (T *)(array + N), val);
}

void dump_func() { DUMPOUT << endl; }
template <class Head, class... Tail>
void dump_func(Head &&head, Tail &&...tail) {
  DUMPOUT << head;
  if (sizeof...(Tail) == 0) {
    DUMPOUT << " ";
  } else {
    DUMPOUT << ", ";
  }
  dump_func(move(tail)...);
}

#define PI 3.14159265358979323846
#define EPS 1e-8
#define FOR(t, a, n) for (int t = (a); t < (n); ++t)
#define REP(t, n) FOR(t, 0, n)
#define all(j) (j).begin(), (j).end()
#define SZ(j) ((int)(j).size())
#define fake false

class Timer {
public:
  double t = 0;
  Timer() {}
  static double time() {
#ifdef _MSC_VER
    return __rdtsc() / 2.8e9;
#else
    unsigned long long a, d;
    __asm__ volatile("rdtsc" : "=a"(a), "=d"(d));
    return (d << 32 | a) / 2.8e9;
#endif
  }
  void measure() { t = time() - t; }
  double elapsedMs() { return (time() - t) * 1000.0; }
} timer;

struct Xorshift {
  uint64_t x = 88172645463325252LL;
  unsigned next_int() {
    x = x ^ (x << 7);
    return x = x ^ (x >> 9);
  }
  unsigned next_int(unsigned mod) {
    x = x ^ (x << 7);
    x = x ^ (x >> 9);
    return x % mod;
  }
  unsigned next_int(unsigned l, unsigned r) {
    x = x ^ (x << 7);
    x = x ^ (x >> 9);
    return x % (r - l + 1) + l;
  }
  double next_double() { return double(next_int()) / UINT_MAX; }
} rnd;

template <typename T> void shuffle_vector(vector<T> &v, Xorshift &rnd) {
  int n = v.size();
  for (int i = n - 1; i >= 1; i--) {
    int r = rnd.next_int(i);
    swap(v[i], v[r]);
  }
}

ll solve(ll N, ll K, vector<ll> A) {

  vector<ll> neg, pos;
  ll zeros = 0;

  REP(i, N) {
    if (A[i] < 0)
      neg.push_back(A[i]);
    else if (A[i] > 0)
      pos.push_back(A[i]);
    else
      zeros++;
  }

  ll below_zero = neg.size() * pos.size();
  ll just_zero = zeros * (N - zeros) + zeros * (zeros - 1) / 2;
  ll upper_zero = N * (N - 1) / 2 - just_zero - below_zero;

  if (below_zero < K && K <= below_zero + just_zero) {
    return 0;
  }

  vector<ll> neg_a(neg), neg_d(neg), pos_a(pos), pos_d(pos);
  sort(all(neg_a));
  sort(all(neg_d), greater<ll>());
  sort(all(pos_a));
  sort(all(pos_d), greater<ll>());

  auto bin = [&](ll ng, ll ok, ll x, ll n, const vector<ll> &v) {
    while (abs(ok - ng) > 1) {
      ll mid = (ng + ok) / 2;
      (n * v[mid] <= x ? ok : ng) = mid;
    }
    return ok;
  };

  auto under_x = [&](ll x) {
    ll ret = 0;
    if (x < 0) {
      if (neg_a.front() * pos_d.front() > x)
        return ret;
      for (ll n : neg_a) {
        if (n * pos_d.front() > x)
          continue;
        if (n * pos_d.back() <= x) {
          ret += pos_d.size();
          continue;
        }
        // n * p <= x なる x の個数を求める
        ret += bin(pos_d.size() - 1, 0, x, n, pos_d) + 1;
      }
    } else {
      ret = below_zero + just_zero;
      // n1 * n2 <= x
      for (int i = 0; i < (int)neg_d.size() - 1; i++) {
        ll n1 = neg_d[i], n2 = neg_d[i + 1];
        if (n1 * n2 > x)
          break;
        if (n1 * neg_d.back() <= x) {
          ret += neg_d.size() - i - 1;
          continue;
        }
        ret += bin(neg_d.size() - 1, i + 1, x, n1, neg_d) - i;
      }
      // p1 * p2 <= x
      for (int i = 0; i < (int)pos_a.size() - 1; i++) {
        ll p1 = pos_a[i], p2 = pos_a[i + 1];
        if (p1 * p2 > x)
          break;
        if (p1 * pos_a.back() <= x) {
          ret += pos_a.size() - i - 1;
          continue;
        }
        ret += bin(pos_a.size() - 1, i + 1, x, p1, pos_a) - i;
      }
    }
    return ret;
  };

  auto under_x_naive = [&](ll x) {
    ll ret = 0;
    for (int i = 0; i < (int)A.size() - 1; i++) {
      for (int j = i + 1; j < A.size(); j++) {
        if (A[i] * A[j] <= x)
          ret++;
      }
    }
    return ret;
  };

  auto bin2 = [&](ll ng, ll ok, ll k) {
    // x 以下の個数が k 以上となるはじめての x を求める
    while (abs(ok - ng) > 1) {
      ll mid = (ng + ok) / 2;
      if (under_x_naive(mid) != under_x(mid)) {
        under_x_naive(mid);
        under_x(mid);
      }
      (under_x_naive(mid) >= k ? ok : ng) = mid;
    }
    return ok;
  };

  return bin2(INT64_MIN / 2, INT64_MAX / 2, K);
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  ll N, K;
  cin >> N >> K;
  vector<ll> A(N);
  cin >> A;

  cout << solve(N, K, A) << endl;

  return 0;
}


Fixed submission:
// #define NDEBUG
#include "bits/stdc++.h"
#include <array>
#include <iostream>
#include <random>
#include <string>
#include <unordered_map>
#include <unordered_set>
#ifdef _MSC_VER
#include <ppl.h>
// #include <boost/multiprecision/cpp_dec_float.hpp>
// #include <boost/multiprecision/cpp_int.hpp>
// #include <boost/rational.hpp>
// #else
// #include <omp.h>
#endif

using namespace std;

#define DUMPOUT cerr
#define dump(...)                                                              \
  DUMPOUT << "  ";                                                             \
  DUMPOUT << #__VA_ARGS__ << " :[" << __LINE__ << ":" << __FUNCTION__ << "]"   \
          << endl;                                                             \
  DUMPOUT << "    ";                                                           \
  dump_func(__VA_ARGS__)

using uint = unsigned;
using ll = long long;
using ull = unsigned long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using pdd = pair<double, double>;
using pss = pair<string, string>;
template <typename _KTy, typename _Ty>
ostream &operator<<(ostream &o, const pair<_KTy, _Ty> &m) {
  o << "{" << m.first << ", " << m.second << "}";
  return o;
}
template <typename _KTy, typename _Ty>
ostream &operator<<(ostream &o, const map<_KTy, _Ty> &m) {
  if (m.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *m.begin();
  for (auto itr = ++m.begin(); itr != m.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _KTy, typename _Ty>
ostream &operator<<(ostream &o, const unordered_map<_KTy, _Ty> &m) {
  if (m.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *m.begin();
  for (auto itr = ++m.begin(); itr != m.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const vector<_Ty> &v) {
  if (v.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << v.front();
  for (auto itr = ++v.begin(); itr != v.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const deque<_Ty> &v) {
  if (v.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << v.front();
  for (auto itr = ++v.begin(); itr != v.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const set<_Ty> &s) {
  if (s.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *(s.begin());
  for (auto itr = ++s.begin(); itr != s.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty>
ostream &operator<<(ostream &o, const unordered_set<_Ty> &s) {
  if (s.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << *(s.begin());
  for (auto itr = ++s.begin(); itr != s.end(); itr++) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const stack<_Ty> &s) {
  if (s.empty()) {
    o << "{ }";
    return o;
  }
  stack<_Ty> t(s);
  o << "{" << t.top();
  t.pop();
  while (!t.empty()) {
    o << ", " << t.top();
    t.pop();
  }
  o << "}";
  return o;
}
template <typename _Ty> ostream &operator<<(ostream &o, const list<_Ty> &l) {
  if (l.empty()) {
    o << "{ }";
    return o;
  }
  o << "{" << l.front();
  for (auto itr = ++l.begin(); itr != l.end(); ++itr) {
    o << ", " << *itr;
  }
  o << "}";
  return o;
}
template <typename _KTy, typename _Ty>
istream &operator>>(istream &is, pair<_KTy, _Ty> &m) {
  is >> m.first >> m.second;
  return is;
}
template <typename _Ty> istream &operator>>(istream &is, vector<_Ty> &v) {
  for (size_t t = 0; t < v.size(); t++)
    is >> v[t];
  return is;
}
template <typename _Ty> istream &operator>>(istream &is, deque<_Ty> &v) {
  for (size_t t = 0; t < v.size(); t++)
    is >> v[t];
  return is;
}
namespace aux { // print tuple
template <typename Ty, unsigned N, unsigned L> struct tp {
  static void print(ostream &os, const Ty &v) {
    os << get<N>(v) << ", ";
    tp<Ty, N + 1, L>::print(os, v);
  }
};
template <typename Ty, unsigned N> struct tp<Ty, N, N> {
  static void print(ostream &os, const Ty &v) { os << get<N>(v); }
};
} // namespace aux

template <typename... Tys>
ostream &operator<<(ostream &os, const tuple<Tys...> &t) {
  os << "{";
  aux::tp<tuple<Tys...>, 0, sizeof...(Tys) - 1>::print(os, t);
  os << "}";
  return os;
}

template <typename A, size_t N, typename T>
inline void Fill(A (&array)[N], const T &val) {
  fill((T *)array, (T *)(array + N), val);
}

void dump_func() { DUMPOUT << endl; }
template <class Head, class... Tail>
void dump_func(Head &&head, Tail &&...tail) {
  DUMPOUT << head;
  if (sizeof...(Tail) == 0) {
    DUMPOUT << " ";
  } else {
    DUMPOUT << ", ";
  }
  dump_func(move(tail)...);
}

#define PI 3.14159265358979323846
#define EPS 1e-8
#define FOR(t, a, n) for (int t = (a); t < (n); ++t)
#define REP(t, n) FOR(t, 0, n)
#define all(j) (j).begin(), (j).end()
#define SZ(j) ((int)(j).size())
#define fake false

class Timer {
public:
  double t = 0;
  Timer() {}
  static double time() {
#ifdef _MSC_VER
    return __rdtsc() / 2.8e9;
#else
    unsigned long long a, d;
    __asm__ volatile("rdtsc" : "=a"(a), "=d"(d));
    return (d << 32 | a) / 2.8e9;
#endif
  }
  void measure() { t = time() - t; }
  double elapsedMs() { return (time() - t) * 1000.0; }
} timer;

struct Xorshift {
  uint64_t x = 88172645463325252LL;
  unsigned next_int() {
    x = x ^ (x << 7);
    return x = x ^ (x >> 9);
  }
  unsigned next_int(unsigned mod) {
    x = x ^ (x << 7);
    x = x ^ (x >> 9);
    return x % mod;
  }
  unsigned next_int(unsigned l, unsigned r) {
    x = x ^ (x << 7);
    x = x ^ (x >> 9);
    return x % (r - l + 1) + l;
  }
  double next_double() { return double(next_int()) / UINT_MAX; }
} rnd;

template <typename T> void shuffle_vector(vector<T> &v, Xorshift &rnd) {
  int n = v.size();
  for (int i = n - 1; i >= 1; i--) {
    int r = rnd.next_int(i);
    swap(v[i], v[r]);
  }
}

ll solve(ll N, ll K, vector<ll> A) {

  vector<ll> neg, pos;
  ll zeros = 0;

  REP(i, N) {
    if (A[i] < 0)
      neg.push_back(A[i]);
    else if (A[i] > 0)
      pos.push_back(A[i]);
    else
      zeros++;
  }

  ll below_zero = neg.size() * pos.size();
  ll just_zero = zeros * (N - zeros) + zeros * (zeros - 1) / 2;
  ll upper_zero = N * (N - 1) / 2 - just_zero - below_zero;

  if (below_zero < K && K <= below_zero + just_zero) {
    return 0;
  }

  vector<ll> neg_a(neg), neg_d(neg), pos_a(pos), pos_d(pos);
  sort(all(neg_a));
  sort(all(neg_d), greater<ll>());
  sort(all(pos_a));
  sort(all(pos_d), greater<ll>());

  auto bin = [&](ll ng, ll ok, ll x, ll n, const vector<ll> &v) {
    while (abs(ok - ng) > 1) {
      ll mid = (ng + ok) / 2;
      (n * v[mid] <= x ? ok : ng) = mid;
    }
    return ok;
  };

  auto under_x = [&](ll x) {
    ll ret = 0;
    if (x < 0) {
      if (neg_a.front() * pos_d.front() > x)
        return ret;
      for (ll n : neg_a) {
        if (n * pos_d.front() > x)
          continue;
        if (n * pos_d.back() <= x) {
          ret += pos_d.size();
          continue;
        }
        // n * p <= x なる x の個数を求める
        ret += bin(pos_d.size() - 1, 0, x, n, pos_d) + 1;
      }
    } else {
      ret = below_zero + just_zero;
      // n1 * n2 <= x
      for (int i = 0; i < (int)neg_d.size() - 1; i++) {
        ll n1 = neg_d[i], n2 = neg_d[i + 1];
        if (n1 * n2 > x)
          break;
        if (n1 * neg_d.back() <= x) {
          ret += neg_d.size() - i - 1;
          continue;
        }
        ret += bin(neg_d.size() - 1, i + 1, x, n1, neg_d) - i;
      }
      // p1 * p2 <= x
      for (int i = 0; i < (int)pos_a.size() - 1; i++) {
        ll p1 = pos_a[i], p2 = pos_a[i + 1];
        if (p1 * p2 > x)
          break;
        if (p1 * pos_a.back() <= x) {
          ret += pos_a.size() - i - 1;
          continue;
        }
        ret += bin(pos_a.size() - 1, i + 1, x, p1, pos_a) - i;
      }
    }
    return ret;
  };

  auto under_x_naive = [&](ll x) {
    ll ret = 0;
    for (int i = 0; i < (int)A.size() - 1; i++) {
      for (int j = i + 1; j < A.size(); j++) {
        if (A[i] * A[j] <= x)
          ret++;
      }
    }
    return ret;
  };

  auto bin2 = [&](ll ng, ll ok, ll k) {
    // x 以下の個数が k 以上となるはじめての x を求める
    while (abs(ok - ng) > 1) {
      ll mid = (ng + ok) / 2;
      (under_x(mid) >= k ? ok : ng) = mid;
    }
    return ok;
  };

  return bin2(INT64_MIN / 2, INT64_MAX / 2, K);
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  ll N, K;
  cin >> N >> K;
  vector<ll> A(N);
  cin >> A;

  cout << solve(N, K, A) << endl;

  return 0;
}