Problem p03117 (C++) — buggy submission, verdict: Time Limit Exceeded

#define _USE_MATH_DEFINES
#include <bits/stdc++.h>
using namespace std;

// template
#define rep(i, a, b) for (int i = (a); i < (b); i++)
#define rrep(i, a, b) for (int i = (a); i > (b); i--)
#define ALL(v) (v).begin(), (v).end()
typedef long long int ll;
const int inf = 0x3fffffff;
const ll INF = 0x1fffffffffffffff;
const double eps = 1e-12;
void tostr(ll x, string &res) {
  while (x)
    res += ('0' + (x % 10)), x /= 10;
  reverse(ALL(res));
  return;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
// end

typedef pair<ll, ll> P;

constexpr ll C = 1000000007;
// 傾きと作る箱の個数から必要なボールの個数
P need(int k, int p) {
  int q = C - p;
  ll lb = 0, rb = C * 2020;
  while (rb - lb > 1) {
    ll mid = (lb + rb) / 2, cnt = 0;
    rep(x, 0, mid / p + 1) {
      cnt += (mid - x * p) / q + 1;
      if (cnt >= k)
        break;
    }
    if (cnt < k)
      lb = mid;
    else
      rb = mid;
  }
  P res({0, 0});
  rep(x, 0, rb / p + 1) {
    ll y = (rb - x * p) / q;
    res.first += x * (y + 1);
    res.second += y * (y + 1) / 2;
  }
  return res;
}

// 箱の個数とボールの個数から判定
bool check(int k, int r, int b) {
  int lb = 0, rb = C + 1;
  P x, y;
  while (rb - lb > 1) {
    int mid = (lb + rb) / 2;
    P p = need(k, mid);
    if (p.first <= r && p.second <= b)
      return true;
    if (p.first > r && p.second > b)
      return false;
    if (p.first > r)
      x = p, lb = mid;
    else
      y = p, rb = mid;
  }
  return (y.first - x.first) * (b - x.second) <=
         (y.second - x.second) * (r - x.first);
}

// reference:https://physics0523.hatenablog.com/entry/2020/02/05/024612
int main() {
  int r, b;
  cin >> r >> b;
  if (r > b)
    swap(r, b);
  int lb = 0, rb = 1341398 * 2;
  while (rb - lb > 1) {
    int mid = (lb + rb) / 2;
    if (check(mid, r, b))
      lb = mid;
    else
      rb = mid;
  }
  cout << lb - 1 << endl;
  return 0;
}


Error: TLE

Fixed submission:

#define _USE_MATH_DEFINES
#include <bits/stdc++.h>
using namespace std;

// template
#define rep(i, a, b) for (int i = (a); i < (b); i++)
#define rrep(i, a, b) for (int i = (a); i > (b); i--)
#define ALL(v) (v).begin(), (v).end()
typedef long long int ll;
const int inf = 0x3fffffff;
const ll INF = 0x1fffffffffffffff;
const double eps = 1e-12;
void tostr(ll x, string &res) {
  while (x)
    res += ('0' + (x % 10)), x /= 10;
  reverse(ALL(res));
  return;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
// end

typedef pair<ll, ll> P;

constexpr int C = 1000000007;
// 傾きと作る箱の個数から必要なボールの個数
P need(int k, int p) {
  int q = C - p;
  ll lb = 0, rb = 2020LL * C;
  while (rb - lb > 1) {
    ll mid = (lb + rb) / 2, cnt = 0;
    rep(x, 0, mid / p + 1) {
      cnt += (mid - x * p) / q + 1;
      if (cnt >= k)
        break;
    }
    if (cnt < k)
      lb = mid;
    else
      rb = mid;
  }
  P res({0, 0});
  rep(x, 0, rb / p + 1) {
    ll y = (rb - x * p) / q;
    res.first += x * (y + 1);
    res.second += y * (y + 1) / 2;
  }
  return res;
}

// 箱の個数とボールの個数から判定
bool check(int k, int r, int b) {
  int lb = 0, rb = C + 1;
  P x, y;
  while (rb - lb > 1) {
    int mid = (lb + rb) / 2;
    P p = need(k, mid);
    if (p.first <= r && p.second <= b)
      return true;
    if (p.first > r && p.second > b)
      return false;
    if (p.first > r)
      x = p, lb = mid;
    else
      y = p, rb = mid;
  }
  return (y.first - x.first) * (b - x.second) <=
         (y.second - x.second) * (r - x.first);
}

// reference:https://physics0523.hatenablog.com/entry/2020/02/05/024612
int main() {
  int r, b;
  cin >> r >> b;
  int lb = 0, rb = 1341398 * 2;
  while (rb - lb > 1) {
    int mid = (lb + rb) / 2;
    if (check(mid, r, b))
      lb = mid;
    else
      rb = mid;
  }
  cout << lb - 1 << endl;
  return 0;
}