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

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

#define pb push_back
#define mp make_pair
#define PI (acos(-1.0))
#define FOR(i, a, b) for (i = a; i <= b; i++)
#define DBG printf("Hi\\n")
#define i64 long long int
#define eps (1e-8)
#define xx first
#define yy second
#define ln 17
#define off 2
#define SZ(z) ((int)z.size())
#define MEM(a, x) memset(a, x, sizeof(a))

#define FastIO                                                                 \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL)

using namespace __gnu_pbds;
using namespace std;

typedef tree<i64, null_type, less<i64>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;

#define IN freopen("262144.in", "r", stdin)
#define OUT freopen("262144.out", "w", stdout)

#define maxn 300005
#define INF 2000000000000000000LL
#define mod 998244353LL
#define log 60

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

typedef pair<i64, i64> pii;

class numberTheory {

public:
  numberTheory() {}

  pii extendedEuclid(i64 a, i64 b) { // returns x, y | ax + by = gcd(a,b)

    if (b == 0)
      return pii(a >= 0 ? 1 : -1, 0LL);
    else {
      pii d = extendedEuclid(b, a % b);
      return pii(d.yy, d.xx - d.yy * (a / b));
    }
  }

  i64 modularInverse(i64 a, i64 n) {
    pair<i64, i64> ret = extendedEuclid(a, n);
    return ((ret.xx % n) + n) % n;
  }

  i64 bigMod(i64 a, i64 n, i64 m) {
    if (n == 0)
      return 1;
    i64 ret = bigMod(a, n / 2, m);
    ret = (ret * ret) % m;
    if (n % 2)
      ret = (ret * a) % m;
    return ret;
  }
};

i64 mulmod(i64 x, i64 y, i64 m) {
  if (x == 0)
    return 0;
  i64 r = (mulmod(x / 2, y, m) * 2) % m;
  if (x % 2)
    r = (r + y) % m;
  return r;
}

i64 gcd(i64 x, i64 y) {
  if (y == 0)
    return x;
  return gcd(y, x % y);
}

void Try(i64 x, i64 y, i64 &ans) {
  if (y % 2 == 0)
    return;
  numberTheory nt;

  i64 xinv = nt.modularInverse(x, y);
  i64 inv2 = nt.modularInverse(2, y);

  i64 n = x * y;

  i64 ret = mulmod(xinv, x, n);
  ret = mulmod(ret, inv2, n);

  i64 ret1 = ret, ret2 = n - ret;
  if (ret1 == 0)
    ret1 = n;
  if (ret2 == 0)
    ret2 = n;

  //    printf("%lld %lld %lld %lld %lld\\n",ret1,x,y, ret1%x , (2*ret1-1)%y ) ;
  //    printf("%lld %lld %lld %lld %lld\\n",ret2,x,y, ret2%x , (2*ret2+1)%y ) ;
  //    printf("\\n") ;

  ans = min(ans, min(2 * ret2, 2 * ret1 - 1));
}

i64 solve(i64 n) {
  i64 ans = INF;

  for (int i = 1; i <= 100; i++) {
    if (((i * (i + 1)) / 2) % n == 0) {
      return i;
    }
  }

  for (int i = 1; 1LL * i * i <= n; i++) {
    if (n % i != 0)
      continue;

    if (gcd(i, n / i) == 1) {
      i64 x = i, y = (n / i);
      Try(x, y, ans);
      Try(y, x, ans);
    }
  }

  if (ans == INF)
    ans = n - 1;
  return ans;
}

i64 brute(i64 n) {
  for (int i = 1;; i++) {
    if (((1LL * i * (i + 1)) / 2) % n == 0)
      return i;
  }
  return n;
}

int main() {
  int tc, t = 1;
  tc = 1;

  while (t++ <= tc) {
    i64 n;
    scanf("%lld", &n);
    //        n = rng()%1000 ;
    //      printf("%lld\\n",n) ;
    //      printf("%lld %lld\\n",brute(n),solve(n)) ;
    printf("%lld\\n", solve(n));
  }
  /*
      while(1)
      {
          i64 n ;
          scanf("%lld",&n) ;
          printf("%lld %lld\\n",brute(n),solve(n)) ;
      }
  */
  //    printf("%lld\\n",ans) ;

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

#define pb push_back
#define mp make_pair
#define PI (acos(-1.0))
#define FOR(i, a, b) for (i = a; i <= b; i++)
#define DBG printf("Hi\\n")
#define i64 long long int
#define eps (1e-8)
#define xx first
#define yy second
#define ln 17
#define off 2
#define SZ(z) ((int)z.size())
#define MEM(a, x) memset(a, x, sizeof(a))

#define FastIO                                                                 \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL)

using namespace __gnu_pbds;
using namespace std;

typedef tree<i64, null_type, less<i64>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;

#define IN freopen("262144.in", "r", stdin)
#define OUT freopen("262144.out", "w", stdout)

#define maxn 300005
#define INF 2000000000000000000LL
#define mod 998244353LL
#define log 60

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

typedef pair<i64, i64> pii;

class numberTheory {

public:
  numberTheory() {}

  pii extendedEuclid(i64 a, i64 b) { // returns x, y | ax + by = gcd(a,b)

    if (b == 0)
      return pii(a >= 0 ? 1 : -1, 0LL);
    else {
      pii d = extendedEuclid(b, a % b);
      return pii(d.yy, d.xx - d.yy * (a / b));
    }
  }

  i64 modularInverse(i64 a, i64 n) {
    pair<i64, i64> ret = extendedEuclid(a, n);
    return ((ret.xx % n) + n) % n;
  }

  i64 bigMod(i64 a, i64 n, i64 m) {
    if (n == 0)
      return 1;
    i64 ret = bigMod(a, n / 2, m);
    ret = (ret * ret) % m;
    if (n % 2)
      ret = (ret * a) % m;
    return ret;
  }
};

i64 mulmod(i64 x, i64 y, i64 m) {
  if (x == 0)
    return 0;
  i64 r = (mulmod(x / 2, y, m) * 2) % m;
  if (x % 2)
    r = (r + y) % m;
  return r;
}

i64 gcd(i64 x, i64 y) {
  if (y == 0)
    return x;
  return gcd(y, x % y);
}

void Try(i64 x, i64 y, i64 &ans) {
  if (y % 2 == 0)
    return;
  numberTheory nt;

  i64 xinv = nt.modularInverse(x, y);
  i64 inv2 = nt.modularInverse(2, y);

  i64 n = x * y;

  i64 ret = mulmod(xinv, x, n);
  ret = mulmod(ret, inv2, n) % mod;

  i64 ret1 = ret, ret2 = n - ret;
  if (ret1 == 0)
    ret1 = n;
  if (ret2 == 0)
    ret2 = n;
  ans = min(ans, min(2 * ret2, 2 * ret1 - 1));
}

i64 solve(i64 n) {
  i64 ans = INF;

  for (int i = 1; i <= 100; i++) {
    if (((i * (i + 1)) / 2) % n == 0) {
      return i;
    }
  }

  for (int i = 1; i * i <= n; i++) {
    if (n % i == 0 && gcd(i, n / i) == 1) {
      i64 x = i, y = (n / i);
      Try(x, y, ans);
      Try(y, x, ans);
    }
  }

  if (ans == INF)
    ans = n - 1;
  return ans;
}

i64 brute(i64 n) {
  for (int i = 1;; i++) {
    if (((i * (i + 1)) / 2) % n == 0)
      return i;
  }
  return n;
}

int main() {
  int tc, t = 1;
  tc = 1;

  while (t++ <= tc) {
    i64 n;
    scanf("%lld", &n);
    //        n = rng()%1000 ;
    //      printf("%lld\\n",n) ;
    //    printf("%lld %lld\\n",brute(n),solve(n)) ;
    printf("%lld\\n", solve(n));
  }
  /*
      while(1)
      {
          i64 n ;
          scanf("%lld",&n) ;
          printf("%lld %lld\\n",brute(n),solve(n)) ;
      }
  */
  //    printf("%lld\\n",ans) ;

  return 0;
}