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

#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
#define ll long long
#define REP(i, a, b) for (ll i = a; i < b; i++)
#define REPI(i, a, b) for (ll i = b - 1; i >= a; i--)
#define i_os                                                                   \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define INF (ll)1e18 + 100
#define endl "\\n"
#define p0(a) cout << a << " "
#define p1(a) cout << a << endl
#define p2(a, b) cout << a << " " << b << endl
#define p3(a, b, c) cout << a << " " << b << " " << c << endl
#define p4(a, b, c, d) cout << a << " " << b << " " << c << " " << d << endl
/*SOME BITMASK KNOWLEDGE
1)x & (x - 1):sets the last one bit of x to zero
power of two exactly when x & (x − 1) = 0.
2)x & -x:sets all the one bits to zero, except last one bit
3)x | (x - 1):inverts all the bits after the last one bit*/
#define o_set                                                                  \\
  tree<int, null_type, less<int>, rb_tree_tag,                                 \\
       tree_order_statistics_node_update>
#define o_setll                                                                \\
  tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update>
typedef tree<pair<ll, ll>, null_type, less<pair<ll, ll>>, rb_tree_tag,
             tree_order_statistics_node_update>
    indexed_set;
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    indexed_set1;
// 1. order_of_key(k) : number of elements strictly lesser than k
// 2. find_by_order(k) : k-th element in the set

int main() {
  i_os;
  ll n;
  cin >> n;
  vector<ll> arr(n + 1);
  vector<bool> prime(n + 1);
  REP(i, 1, n + 1) {
    arr[i] = i;
    prime[i] = 1;
  }
  REP(i, 2, sqrt(n) + 1) {
    if (prime[i]) {
      for (ll j = i; j <= n / i; j++) {
        prime[i * j] = 0;
        if (arr[i * j] == i * j) {
          arr[i * j] = i;
        }
      }
    }
  }

  vector<ll> ans(n + 1);
  REP(i, 1, n + 1) ans[i] = 1;
  REP(i, 2, n + 1) {
    ll val2 = i;
    ll ct = 0;
    while (val2 % arr[i] == 0) {
      val2 /= arr[i];
      ct++;
    }
    ans[i] = (ct + 1) * ans[val2];
  }
  ll p = 0;
  REP(i, 1, n + 1) { p += i * ans[i]; }
  cout << p << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
#define ll long long
#define REP(i, a, b) for (ll i = a; i < b; i++)
#define REPI(i, a, b) for (ll i = b - 1; i >= a; i--)
#define i_os                                                                   \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define INF (ll)1e18 + 100
#define endl "\\n"
#define p0(a) cout << a << " "
#define p1(a) cout << a << endl
#define p2(a, b) cout << a << " " << b << endl
#define p3(a, b, c) cout << a << " " << b << " " << c << endl
#define p4(a, b, c, d) cout << a << " " << b << " " << c << " " << d << endl
/*SOME BITMASK KNOWLEDGE
1)x & (x - 1):sets the last one bit of x to zero
power of two exactly when x & (x − 1) = 0.
2)x & -x:sets all the one bits to zero, except last one bit
3)x | (x - 1):inverts all the bits after the last one bit*/
#define o_set                                                                  \\
  tree<int, null_type, less<int>, rb_tree_tag,                                 \\
       tree_order_statistics_node_update>
#define o_setll                                                                \\
  tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update>
typedef tree<pair<ll, ll>, null_type, less<pair<ll, ll>>, rb_tree_tag,
             tree_order_statistics_node_update>
    indexed_set;
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    indexed_set1;
// 1. order_of_key(k) : number of elements strictly lesser than k
// 2. find_by_order(k) : k-th element in the set

int main() {
  ll n;
  cin >> n;
  vector<ll> arr(n + 1);
  vector<bool> prime(n + 1);
  REP(i, 1, n + 1) {
    arr[i] = i;
    prime[i] = 1;
  }
  REP(i, 2, sqrt(n) + 1) {
    if (prime[i]) {
      for (ll j = 2; j <= n / i; j++) {
        prime[i * j] = 1;
        if (arr[i * j] == i * j) {
          arr[i * j] = i;
        }
      }
    }
  }
  vector<ll> ans(n + 1);
  REP(i, 1, n + 1) ans[i] = 1;
  REP(i, 2, n + 1) {
    ll val2 = i;
    ll ct = 0;
    while (val2 % arr[i] == 0) {
      val2 /= arr[i];
      ct++;
    }
    ans[i] = (ct + 1) * ans[val2];
  }
  ll p = 0;
  REP(i, 1, n + 1) { p += i * ans[i]; }
  cout << p << endl;
  return 0;
}