Problem p01550 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
ll mod = 1e9 + 7;

ll add(ll a, ll b) { return (a + b) % mod; }

void Add(int &a, ll b) {
  ll tmp = (a + b) % mod;
  a = tmp;
}

ll mul(ll a, ll b) { return (a * b) % mod; }

ll mpow(ll a, ll b) {
  ll res = 1;
  while (b) {
    if (b % 2 == 1)
      res = mul(res, a);
    a = mul(a, a);
    b /= 2;
  }
  return res;
}

ll divi(ll a, ll b) { return mul(a, mpow(b, mod - 2)); }

int n;
int a[200];

ll Z = 0;
ll cnt[10];
ll cntinv[10];

ll sum[10];

// ll dp[50][50][50][50][2];
// ll DP[50][50][50][50][2];
vector<vector<vector<vector<vector<int>>>>> dp, DP;

void update(int a, int b, int c, int d, int z) {

  int na = a % 2;
  int ma = 1 - na;

  /*  cout<<a<<' '<<b<<' '<<c<<' '<<d<<' '<<z<<endl;
  cout<<dp[a][b][c][d][z]<<endl;
  cout<<DP[a][b][c][d][z]<<endl;
  cout<<endl;    */
  if (a < cnt[1]) {
    int cc = cnt[1] - a;
    ll pp = mul(dp[na][b][c][d][z], cc);
    Add(dp[ma][b][c][d][z], pp);
    Add(DP[ma][b][c][d][z], mul(DP[na][b][c][d][z], 10 * cc));
    Add(DP[ma][b][c][d][z], mul(mul(pp, cntinv[1]), sum[1]));
  }

  if (b < cnt[2]) {
    int cc = cnt[2] - b;
    ll pp = mul(dp[na][b][c][d][z], cc);
    Add(dp[na][b + 1][c][d][z], pp);
    Add(DP[na][b + 1][c][d][z], mul(DP[na][b][c][d][z], 100 * cc));
    Add(DP[na][b + 1][c][d][z], mul(mul(pp, cntinv[2]), sum[2]));
  }
  if (c < cnt[3]) {
    int cc = cnt[3] - c;
    ll pp = mul(dp[na][b][c][d][z], cc);
    Add(dp[na][b][c + 1][d][z], mul(dp[na][b][c][d][z], cc));
    Add(DP[na][b][c + 1][d][z], mul(DP[na][b][c][d][z], 1000 * cc));
    Add(DP[na][b][c + 1][d][z], mul(mul(pp, cntinv[3]), sum[3]));
  }
  if (d < cnt[4]) {
    int cc = cnt[4] - d;
    ll pp = mul(dp[na][b][c][d][z], cc);
    Add(dp[na][b][c][d + 1][z], mul(dp[na][b][c][d][z], cc));
    Add(DP[na][b][c][d + 1][z], mul(DP[na][b][c][d][z], 10000 * cc));
    Add(DP[na][b][c][d + 1][z], mul(mul(pp, cntinv[4]), sum[4]));
  }

  if (z < Z && a + b + c + d > 0) {
    Add(dp[na][b][c][d][z + 1], dp[na][b][c][d][z]);
    Add(DP[na][b][c][d][z + 1], mul(DP[na][b][c][d][z], 10));
  }
}

ll solve() {
  dp.resize(2);
  for (int i = 0; i < (int)dp.size(); i++) {
    dp[i].resize(cnt[2] + 1);
    for (int j = 0; j < (int)dp[i].size(); j++) {
      dp[i][j].resize(cnt[3] + 1);
      for (int k = 0; k < (int)dp[i][j].size(); k++) {
        dp[i][j][k].resize(cnt[4] + 1);
        for (int l = 0; l < (int)dp[i][j][k].size(); l++) {
          dp[i][j][k][l].resize(2);
        }
      }
    }
  }
  DP = dp;

  dp[0][0][0][0][0] = 1;

  int ans = 0;

  for (int a = 0; a <= cnt[1]; a++) {
    for (int b = 0; b <= cnt[2]; b++)
      for (int c = 0; c <= cnt[3]; c++)
        for (int d = 0; d <= cnt[4]; d++)
          for (int z = 0; z <= Z; z++) {
            update(a, b, c, d, z);
            Add(ans, DP[a % 2][b][c][d][z]);

            dp[a % 2][b][c][d][z] = 0;
            DP[a % 2][b][c][d][z] = 0;
          }
  }

  return ans;
}

int countDisit(int x) {
  int res = 0;
  while (x) {
    res++;
    x /= 10;
  }
  return res;
}

int main() {
  cin >> n;

  for (int i = 0; i < n; i++) {
    cin >> a[i];
    int D = countDisit(a[i]);
    if (a[i] == 0) {
      Z++;
      continue;
    }
    cnt[D]++;
    sum[D] += a[i];
  }

  for (int i = 1; i <= 4; i++)
    if (cnt[i] > 0)
      cntinv[i] = divi(1, cnt[i]);
  cout << solve() << endl;

  return 0;
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
ll mod = 1e9 + 7;

ll add(ll a, ll b) { return (a + b) % mod; }

void Add(ll &a, ll b) { a = (a + b) % mod; }

ll mul(ll a, ll b) { return (a * b) % mod; }

ll mpow(ll a, ll b) {
  ll res = 1;
  while (b) {
    if (b % 2 == 1)
      res = mul(res, a);
    a = mul(a, a);
    b /= 2;
  }
  return res;
}

ll divi(ll a, ll b) { return mul(a, mpow(b, mod - 2)); }

ll mem[1000][1000];

ll nCr(ll n, ll r) {
  if (r == 0)
    return 1;
  if (n == r)
    return 1;
  if (n < r)
    return 0;

  if (mem[n][r])
    return mem[n][r];
  mem[n][r] = add(nCr(n - 1, r - 1), nCr(n - 1, r));
  return mem[n][r];
}

int n;
int a[200];

ll Z = 0;
ll cnt[10];
ll cntinv[10];

ll sum[10];

// ll dp[50][50][50][50][2];
// ll DP[50][50][50][50][2];
vector<vector<vector<vector<vector<ll>>>>> dp, DP;

void update(int a, int b, int c, int d, int z) {
  /*  cout<<a<<' '<<b<<' '<<c<<' '<<d<<' '<<z<<endl;
  cout<<dp[a][b][c][d][z]<<endl;
  cout<<DP[a][b][c][d][z]<<endl;
  cout<<endl;    */
  if (a < cnt[1]) {
    int cc = cnt[1] - a;
    ll pp = mul(dp[a][b][c][d][z], cc);
    Add(dp[a + 1][b][c][d][z], pp);
    Add(DP[a + 1][b][c][d][z], mul(DP[a][b][c][d][z], 10 * cc));
    Add(DP[a + 1][b][c][d][z], mul(mul(pp, cntinv[1]), sum[1]));
  }

  if (b < cnt[2]) {
    int cc = cnt[2] - b;
    ll pp = mul(dp[a][b][c][d][z], cc);
    Add(dp[a][b + 1][c][d][z], pp);
    Add(DP[a][b + 1][c][d][z], mul(DP[a][b][c][d][z], 100 * cc));
    Add(DP[a][b + 1][c][d][z], mul(mul(pp, cntinv[2]), sum[2]));
  }
  if (c < cnt[3]) {
    int cc = cnt[3] - c;
    ll pp = mul(dp[a][b][c][d][z], cc);
    Add(dp[a][b][c + 1][d][z], mul(dp[a][b][c][d][z], cc));
    Add(DP[a][b][c + 1][d][z], mul(DP[a][b][c][d][z], 1000 * cc));
    Add(DP[a][b][c + 1][d][z], mul(mul(pp, cntinv[3]), sum[3]));
  }
  if (d < cnt[4]) {
    int cc = cnt[4] - d;
    ll pp = mul(dp[a][b][c][d][z], cc);
    Add(dp[a][b][c][d + 1][z], mul(dp[a][b][c][d][z], cc));
    Add(DP[a][b][c][d + 1][z], mul(DP[a][b][c][d][z], 10000 * cc));
    Add(DP[a][b][c][d + 1][z], mul(mul(pp, cntinv[4]), sum[4]));
  }

  if (z < Z && a + b + c + d > 0) {
    Add(dp[a][b][c][d][z + 1], dp[a][b][c][d][z]);
    Add(DP[a][b][c][d][z + 1], mul(DP[a][b][c][d][z], 10));
  }
}

ll solve() {
  dp.resize(cnt[1] + 1);
  for (int i = 0; i < (int)dp.size(); i++) {
    dp[i].resize(cnt[2] + 1);
    for (int j = 0; j < (int)dp[i].size(); j++) {
      dp[i][j].resize(cnt[3] + 1);
      for (int k = 0; k < (int)dp[i][j].size(); k++) {
        dp[i][j][k].resize(cnt[4] + 1);
        for (int l = 0; l < (int)dp[i][j][k].size(); l++) {
          dp[i][j][k][l].resize(2);
        }
      }
    }
  }
  DP = dp;

  dp[0][0][0][0][0] = 1;

  ll ans = 0;

  for (int a = 0; a <= cnt[1]; a++)
    for (int b = 0; b <= cnt[2]; b++)
      for (int c = 0; c <= cnt[3]; c++)
        for (int d = 0; d <= cnt[4]; d++)
          for (int z = 0; z <= Z; z++) {
            update(a, b, c, d, z);
            Add(ans, DP[a][b][c][d][z]);
          }

  return ans;
}

int countDisit(int x) {
  int res = 0;
  while (x) {
    res++;
    x /= 10;
  }
  return res;
}

int main() {
  cin >> n;

  for (int i = 0; i < n; i++) {
    cin >> a[i];
    int D = countDisit(a[i]);
    if (a[i] == 0) {
      Z++;
      continue;
    }
    cnt[D]++;
    sum[D] += a[i];
  }

  for (int i = 1; i <= 4; i++)
    if (cnt[i] > 0)
      cntinv[i] = divi(1, cnt[i]);
  cout << solve() << endl;

  return 0;
}