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

#include <algorithm>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <utility>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<int, int> PII;
typedef vector<int> VI;
typedef vector<VI> VVI;
#define MP make_pair
#define PB push_back
const ll mod = 1000000007;

ll dp[101][4010];
long long ruijo(long long x, ll n) {
  long long p, ans;
  if (n == 0)
    return 1;
  if (n == 1)
    return x;
  if (n % 2) {
    p = ruijo(x, n / 2);
    ans = (x * p) % mod;
    ans = (ans * p) % mod;
    return ans;
  } else {
    p = ruijo(x, n / 2);
    ans = (p * p) % mod;
    return ans;
  }
}
ll n, k;

ll element_count(ll x, ll y, ll z) {
  ll a = max((ll)0, y * z - x);
  ll b = min(y, y * (z + 1) - 1 - x);
  if (a > b)
    return (ll)0;
  else
    return b - a + 1;
}

int main() {
  cin >> n >> k;
  for (ll i = n; i > k; i--) {
    dp[i][0] = ruijo(k + 1, (n - i + 1));
  }
  dp[k][0] = k * ruijo(k + 1, n - k);
  dp[k][1] = 1 * ruijo(k + 1, n - k);
  for (ll i = k - 1; i >= 1; i--) {
    ll p;
    if (i == 1) {
      p = 4000;
    }
    if (i == 2) {
      p = 3000;
    }
    if (i == 3) {
      p = 2000;
    }
    if (i >= 4) {
      p = 1000;
    }
    for (ll j = 0; j < p; j++) {
      if (j == 0) {
        dp[i][0] = k * dp[i + 1][0] % mod;
      } else {
        for (ll m = 0; m <= j; m++) {
          if (j != m) {
            dp[i][j] =
                (dp[i][j] + element_count(m, i, j - m) * dp[i + 1][m]) % mod;
          } else {
            dp[i][j] = (dp[i][j] + (k - i) * dp[i + 1][m]) % mod;
            if (i - m - 1 >= 0) {
              dp[i][j] = (dp[i][j] + (i - m) * dp[i + 1][m]) % mod;
            }
          }
        }
      }
    }
  }
  ll ans = 0;
  for (int i = 1; i < 4000; i++) {
    ans = (ans + i * dp[1][i]) % mod;
  }
  ll w = (n * ruijo(k + 1, n - 1)) % mod;
  ll s = ((k * (k + 1)) / 2) % mod;
  s = (s * w) % mod;
  cout << (s - ans + mod) % mod << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <utility>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<int, int> PII;
typedef vector<int> VI;
typedef vector<VI> VVI;
#define MP make_pair
#define PB push_back
const ll mod = 1000000007;

ll dp[101][5000];
long long ruijo(long long x, ll n) {
  long long p, ans;
  if (n == 0)
    return 1;
  if (n == 1)
    return x;
  if (n % 2) {
    p = ruijo(x, n / 2);
    ans = (x * p) % mod;
    ans = (ans * p) % mod;
    return ans;
  } else {
    p = ruijo(x, n / 2);
    ans = (p * p) % mod;
    return ans;
  }
}
ll n, k;

ll element_count(ll x, ll y, ll z) {
  ll a = max((ll)0, y * z - x);
  ll b = min(y, y * (z + 1) - 1 - x);
  if (a > b)
    return (ll)0;
  else
    return b - a + 1;
}

int main() {
  cin >> n >> k;
  for (ll i = n; i > k; i--) {
    dp[i][0] = ruijo(k + 1, (n - i + 1));
  }
  dp[k][0] = k * ruijo(k + 1, n - k);
  dp[k][1] = 1 * ruijo(k + 1, n - k);
  for (ll i = k - 1; i >= 1; i--) {
    for (ll j = 0; j < 4000; j++) {
      if (j == 0) {
        dp[i][0] = k * dp[i + 1][0] % mod;
      } else {
        for (ll m = 0; m <= j; m++) {
          if (j != m) {
            dp[i][j] =
                (dp[i][j] + element_count(m, i, j - m) * dp[i + 1][m]) % mod;
          } else {
            dp[i][j] = (dp[i][j] + (k - i) * dp[i + 1][m]) % mod;
            if (i - m - 1 >= 0) {
              dp[i][j] = (dp[i][j] + (i - m) * dp[i + 1][m]) % mod;
            }
          }
        }
      }
    }
  }
  ll ans = 0;
  for (int i = 1; i < 4000; i++) {
    ans = (ans + i * dp[1][i]) % mod;
  }
  ll w = (n * ruijo(k + 1, n - 1)) % mod;
  ll s = ((k * (k + 1)) / 2) % mod;
  s = (s * w) % mod;
  cout << (s - ans + mod) % mod << endl;
  return 0;
}