Problem p02792 (C++)

Buggy submission (status: Time Limit Exceeded):
#pragma warning(disable : 4996)
#pragma comment(linker, "/STACK:336777216")
#pragma GCC optimize("Ofast")

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

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;
#define int long long int
typedef complex<long double> cd;
const long double pi = acos(-1);
typedef double db;
typedef long double ldb;
typedef pair<int, int> pii;
typedef pair<db, db> pdd;
typedef vector<int> vi;
typedef vector<vector<int>> matrix;

#define m1 make_pair
#define pb push_back
#define flush fflush(stdout)
#define IOS                                                                    \
  std::ios::sync_with_stdio(false);                                            \
  cin.tie(NULL);                                                               \
  cout.tie(NULL);
#define reset(x, v) memset(x, v, sizeof(x))
#define sz(x) (int)(x.size())
#define fi first
#define se second
#define endl "\n"
#define debug(x) (cerr << #x << ": " << x << "\n")
#define setbits(x) __builtin_popcount(x)
#define all(x) x.begin(), x.end()
#define pii pair<int, int>

inline ll gcd(ll a, ll b) {
  if (a == 0)
    return b;
  return gcd(b % a, a);
}
inline ll power(ll a, ll n, ll m) {
  if (n == 0)
    return 1;
  ll p = power(a, n / 2, m);
  p = (p % m * p % m) % m;
  if (n % 2)
    return (p % m * a % m) % m;
  else
    return p;
}
const ll MOD = 998244353;
const ll hell = 1000000007;
const int INF = 1e18;
const ll LL_INF = 0x3f3f3f3f3f3f3f3f;
typedef tree<int, null_type, less<int>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;
const int N = 200005;
int pr[11];
int cnt;
void recur(int x, int l, int dg, int n) {
  if (dg == 1) {
    x = x * 10 + l;
    if (x <= n)
      cnt++;
    return;
  }
  for (int i = 0; i < 10; i++) {
    recur(x * 10 + i, l, dg - 1, n);
  }
}
int count(int f, int l, int n) {
  // int cnt=0;
  cnt = 0;
  if (f == l && f <= n)
    cnt++;
  int dg = log10(n);
  dg++;
  if (dg >= 2) {
    for (int i = 2; i < dg; i++) {
      cnt = cnt + pr[i - 2];
    }
    recur(f, l, dg - 1, n);
  }
  return cnt;
}
void solve() {
  int n;
  cin >> n;

  pr[0] = 1;
  for (int i = 1; i <= 10; i++)
    pr[i] = 10 * pr[i - 1];
  int ans = 0;
  for (int i = 1; i <= n; i++) {
    int last = i % 10;
    int z = log10(i);
    int x = i / pr[z];
    int first = x % 10;
    if (last) {
      int lst = first;
      int fst = last;
      ans = ans + count(fst, lst, n);
    }
  }
  cout << ans << endl;
}

int32_t main() {
  IOS;
  int t;
  t = 1;
  // cin>>t;
  while (t--) {
    solve();
  }
  return 0;
}


Fixed submission:
#pragma warning(disable : 4996)
#pragma comment(linker, "/STACK:336777216")
#pragma GCC optimize("Ofast")

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

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;
#define int long long int
typedef complex<long double> cd;
const long double pi = acos(-1);
typedef double db;
typedef long double ldb;
typedef pair<int, int> pii;
typedef pair<db, db> pdd;
typedef vector<int> vi;
typedef vector<vector<int>> matrix;

#define m1 make_pair
#define pb push_back
#define flush fflush(stdout)
#define IOS                                                                    \
  std::ios::sync_with_stdio(false);                                            \
  cin.tie(NULL);                                                               \
  cout.tie(NULL);
#define reset(x, v) memset(x, v, sizeof(x))
#define sz(x) (int)(x.size())
#define fi first
#define se second
#define endl "\n"
#define debug(x) (cerr << #x << ": " << x << "\n")
#define setbits(x) __builtin_popcount(x)
#define all(x) x.begin(), x.end()
#define pii pair<int, int>

inline ll gcd(ll a, ll b) {
  if (a == 0)
    return b;
  return gcd(b % a, a);
}
inline ll power(ll a, ll n, ll m) {
  if (n == 0)
    return 1;
  ll p = power(a, n / 2, m);
  p = (p % m * p % m) % m;
  if (n % 2)
    return (p % m * a % m) % m;
  else
    return p;
}
const ll MOD = 998244353;
const ll hell = 1000000007;
const int INF = 1e18;
const ll LL_INF = 0x3f3f3f3f3f3f3f3f;
typedef tree<int, null_type, less<int>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;
const int N = 200005;
int pr[11];
int cnt;
void recur(int x, int l, int dg, int n) {
  if (dg == 1) {
    x = x * 10 + l;
    if (x <= n)
      cnt++;
    return;
  }
  for (int i = 0; i < 10; i++) {
    recur(x * 10 + i, l, dg - 1, n);
  }
}
int count(int f, int l, int n) {
  // int cnt=0;
  cnt = 0;
  if (f == l && f <= n)
    cnt++;
  int dg = log10(n);
  dg++;
  if (dg >= 2) {
    for (int i = 2; i < dg; i++) {
      cnt = cnt + pr[i - 2];
    }
    int v = n / pr[dg - 1];
    v = v % 10;
    if (f == v)
      recur(f, l, dg - 1, n);
    else if (f < v) {
      cnt = cnt + pr[dg - 2];
    }
  }
  return cnt;
}
void solve() {
  int n;
  cin >> n;

  pr[0] = 1;
  for (int i = 1; i <= 10; i++)
    pr[i] = 10 * pr[i - 1];
  int ans = 0;
  for (int i = 1; i <= n; i++) {
    int last = i % 10;
    int z = log10(i);
    int x = i / pr[z];
    int first = x % 10;
    if (last) {
      int lst = first;
      int fst = last;
      ans = ans + count(fst, lst, n);
    }
  }
  cout << ans << endl;
}

int32_t main() {
  IOS;
  int t;
  t = 1;
  // cin>>t;
  while (t--) {
    solve();
  }
  return 0;
}