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

#include <algorithm>
#include <bits/stdc++.h>
#include <cmath>
#include <iostream>
#include <string>
#include <unordered_map>
#include <unordered_set>
using namespace std;
#define ll long long
#define rep(i, n) for (ll i = 0; i < (n); i++)
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define FORR(i, a, b) for (ll i = (a); i <= (b); i++)
#define repR(i, n) for (ll i = n; i >= 0; i--)
#define all(v) (v).begin(), (v).end()
#define rall(v) (v).rbegin(), (v).rend()
#define F first
#define S second
#define pb push_back
#define pu push
#define COUT(x) cout << (x) << "\\n"
#define PQ priority_queue<ll>
#define PQR priority_queue<ll, vector<ll>, greater<ll>>
#define YES(n) cout << ((n) ? "YES\\n" : "NO\\n")
#define Yes(n) cout << ((n) ? "Yes\\n" : "No\\n")
#define mp make_pair
#define sz(x) (ll)(x).size()
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
const ll MOD = 1000000007LL;
const ll INF = 1LL << 60;
using vll = vector<ll>;
using vb = vector<bool>;
using vvb = vector<vb>;
using vvll = vector<vll>;
using vstr = vector<string>;
using vc = vector<char>;
using vvc = vector<vc>;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
ll dx[4] = {0, 1, 0, -1};
ll dy[4] = {1, 0, -1, 0};
int main() {
  ll n;
  cin >> n;
  vll a(n);
  rep(i, n) {
    ll cu = i;
    while (cu < n) {
      a[cu]++;
      cu += i + 1;
    }
  }
  ll ans = 0;
  rep(i, n) { ans += a[i] * (i + 1); }
  COUT(ans);
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bits/stdc++.h>
#include <cmath>
#include <iostream>
#include <string>
#include <unordered_map>
#include <unordered_set>
using namespace std;
#define ll long long
#define rep(i, n) for (ll i = 0; i < (n); i++)
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define FORR(i, a, b) for (ll i = (a); i <= (b); i++)
#define repR(i, n) for (ll i = n; i >= 0; i--)
#define all(v) (v).begin(), (v).end()
#define rall(v) (v).rbegin(), (v).rend()
#define F first
#define S second
#define pb push_back
#define pu push
#define COUT(x) cout << (x) << "\\n"
#define PQ priority_queue<ll>
#define PQR priority_queue<ll, vector<ll>, greater<ll>>
#define YES(n) cout << ((n) ? "YES\\n" : "NO\\n")
#define Yes(n) cout << ((n) ? "Yes\\n" : "No\\n")
#define mp make_pair
#define sz(x) (ll)(x).size()
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
const ll MOD = 1000000007LL;
const ll INF = 1LL << 60;
using vll = vector<ll>;
using vb = vector<bool>;
using vvb = vector<vb>;
using vvll = vector<vll>;
using vstr = vector<string>;
using vc = vector<char>;
using vvc = vector<vc>;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
ll dx[4] = {0, 1, 0, -1};
ll dy[4] = {1, 0, -1, 0};
vll arr(10000001);
ll n;
vector<ll> k(10000001);
void Eratosthenes() {
  for (ll i = 2; i < n; i++) {
    arr[i] = 1;
  }
  for (ll i = 2; i < sqrt(n); i++) {
    if (arr[i]) {
      for (ll j = 0; i * (j + 2) < n; j++) {
        arr[i * (j + 2)] = 0;
        if (k[i * (j + 2)] == 0)
          k[i * (j + 2)] = i;
      }
    }
  }

  for (int i = 2; i < n; i++) {
    if (arr[i]) {
      // cout << i << endl;
    }
  }
}
int main() {
  cin >> n;
  n++;
  Eratosthenes();
  vector<map<ll, ll>> p(n);
  ll ans = 1;
  rep(i, n) {
    if (i == 0 || i == 1)
      continue;
    if (arr[i]) {
      p[i][i]++;
      ans += 2 * i;
    } else {
      ll h = i / k[i];
      for (auto e : p[k[i]]) {
        p[i][e.F] += e.S;
      }
      for (auto e : p[h]) {
        p[i][e.F] += e.S;
      }
      ll cu = 1;
      for (auto e : p[i]) {
        cu *= e.S + 1;
      }

      cu *= i;
      ans += cu;
    }
  }
  COUT(ans);
}