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

#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <numeric>
#include <utility>
#include <vector>

constexpr intmax_t mod = 1e9 + 7;

template <class Tp> Tp gcd(Tp a, Tp b, Tp &x, Tp &y) {
  x = Tp(0);
  y = Tp(1);
  for (Tp u = y, v = x; a;) {
    Tp q = b / a;
    std::swap(x -= q * u, u);
    std::swap(y -= q * v, v);
    std::swap(b -= q * a, a);
  }
  return b;
}

template <class Tp> Tp modinv(Tp a, Tp mod) {
  Tp x, y;
  gcd(a, mod, x, y);
  x %= mod;
  if (x < 0)
    x += mod;
  return x;
}

template <class Tp> Tp modadd(Tp a, Tp b, Tp mod) {
  a += b % mod;
  if (a < 0)
    a += mod;
  if (a >= mod)
    a -= mod;
  return a;
}

template <class Tp> Tp modadd(const std::initializer_list<Tp> &adds, Tp mod) {
  Tp res = 0;
  for (const auto &add : adds) {
    res += add % mod;
    if (res < 0)
      res += mod;
    if (res >= mod)
      res -= mod;
  }
  return res;
}

template <class Tp> Tp modsub(Tp a, Tp b, Tp mod) {
  a -= b % mod;
  if (a < 0)
    a += mod;
  if (a >= mod)
    a -= mod;
  return a;
}

template <class Tp> Tp modmul(const std::initializer_list<Tp> &muls, Tp mod) {
  Tp res = 1;
  for (const auto &mul : muls)
    (res *= mul) %= mod;
  return res;
}

template <class Tp> Tp modpow(Tp base, intmax_t iexp, Tp mod) {
  Tp res = 1;
  for (Tp dbl = base; iexp; iexp >>= 1) {
    if (iexp & 1)
      res = res * dbl % mod;
    dbl = dbl * dbl % mod;
  }
  return res;
}

class modfactorial {
  std::vector<intmax_t> fact, fact_inv;
  intmax_t mod;

public:
  modfactorial(intmax_t N, intmax_t mod) : mod(mod) {
    fact.resize(N + 1);
    fact_inv.resize(N + 1);
    fact[0] = 1;
    for (intmax_t i = 1; i <= N; ++i)
      fact[i] = fact[i - 1] * i % mod;

    fact_inv[N] = modinv(fact[N], mod);
    for (intmax_t i = N; i--;)
      fact_inv[i] = fact_inv[i + 1] * (i + 1) % mod;
  }

  intmax_t operator()(intmax_t k) const { return fact[k]; }

  intmax_t inverse(intmax_t k) const { return fact_inv[k]; }
};

template <class Tp> std::vector<Tp> make_vector(size_t n, const Tp &x) {
  return std::vector<Tp>(n, x);
}

template <class... Ts> auto make_vector(size_t n, size_t m, const Ts &...x) {
  return std::vector<decltype(make_vector(m, x...))>(n, make_vector(m, x...));
}

std::vector<std::vector<intmax_t>> make_dp(std::vector<size_t> dig, size_t d) {
  --dig[d];
  size_t n = std::accumulate(dig.begin(), dig.end(), 0);
  size_t m = 0;
  for (size_t i = 1; i <= 4; ++i)
    m += i * dig[i];
  auto dp = make_vector(n + 1, n + 1, m + 1, intmax_t(0));
  dp[0][0][0] = 1;

  std::vector<size_t> ds;
  ds.reserve(n);
  for (size_t i = 1; i <= 4; ++i)
    for (size_t j = 0; j < dig[i]; ++j)
      ds.push_back(i);

  for (size_t i = 0; i < n; ++i) {
    for (size_t j = 0; j < n; ++j) {
      for (size_t k = 0; k <= m; ++k) {
        (dp[i + 1][j][k] += dp[i][j][k]) %= mod;
        if (k + ds[i] <= m) {
          (dp[i + 1][j + 1][k + ds[i]] += (j + 1) * dp[i][j][k]) %= mod;
        }
      }
    }
  }

  return dp[n];
}

intmax_t solve(const std::vector<intmax_t> &a) {
  size_t n = a.size();
  std::vector<size_t> dig(5);
  std::vector<intmax_t> sum(5);
  for (auto ai : a) {
    if (ai == 0) {
      ++dig[1];
      continue;
    }
    int d = 0;
    for (int i = ai; i > 0; i /= 10)
      ++d;
    ++dig[d];
    (sum[d] += ai) %= mod;
  }

  modfactorial mf(n + 1, mod);
  auto coes = make_vector(n, n, intmax_t(0));
  for (size_t j = 0; j < n; ++j)
    for (size_t l = 0; l + j < n; ++l)
      coes[j][l] = mf(n - 1 - j) * mf.inverse(n - 1 - j - l) % mod;

  intmax_t res = 0;
  for (int i = 1; i <= 4; ++i) {
    if (!dig[i])
      continue;
    auto dp = make_dp(dig, i);

    for (size_t j = 0; j < n; ++j) {
      for (size_t l = 0; l + j < n; ++l) {
        intmax_t cur = 0;
        for (size_t k = 0; k < dp[j].size(); ++k) {
          cur += modmul({sum[i], modpow(10L, k, mod), dp[j][k]}, mod);
          if (cur >= mod)
            cur -= mod;
        }
        (res += cur * coes[j][l]) %= mod;
      }
    }
  }
  return res;
}

int main() {
  size_t n;
  scanf("%zu", &n);

  std::vector<intmax_t> a(n);
  for (auto &ai : a)
    scanf("%jd", &ai);

  std::sort(a.rbegin(), a.rend());
  intmax_t res = solve(a);
  if (a[n - 1] == 0) {
    a.pop_back();
    res -= solve(a);
    if (res < 0)
      res += mod;
  }

  printf("%jd\\n", res);
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <numeric>
#include <utility>
#include <vector>

constexpr intmax_t mod = 1e9 + 7;

template <class Tp> Tp gcd(Tp a, Tp b, Tp &x, Tp &y) {
  x = Tp(0);
  y = Tp(1);
  for (Tp u = y, v = x; a;) {
    Tp q = b / a;
    std::swap(x -= q * u, u);
    std::swap(y -= q * v, v);
    std::swap(b -= q * a, a);
  }
  return b;
}

template <class Tp> Tp modinv(Tp a, Tp mod) {
  Tp x, y;
  gcd(a, mod, x, y);
  x %= mod;
  if (x < 0)
    x += mod;
  return x;
}

template <class Tp> Tp modadd(Tp a, Tp b, Tp mod) {
  a += b % mod;
  if (a < 0)
    a += mod;
  if (a >= mod)
    a -= mod;
  return a;
}

template <class Tp> Tp modadd(const std::initializer_list<Tp> &adds, Tp mod) {
  Tp res = 0;
  for (const auto &add : adds) {
    res += add % mod;
    if (res < 0)
      res += mod;
    if (res >= mod)
      res -= mod;
  }
  return res;
}

template <class Tp> Tp modsub(Tp a, Tp b, Tp mod) {
  a -= b % mod;
  if (a < 0)
    a += mod;
  if (a >= mod)
    a -= mod;
  return a;
}

template <class Tp> Tp modmul(const std::initializer_list<Tp> &muls, Tp mod) {
  Tp res = 1;
  for (const auto &mul : muls)
    (res *= mul) %= mod;
  return res;
}

template <class Tp> Tp modpow(Tp base, intmax_t iexp, Tp mod) {
  Tp res = 1;
  for (Tp dbl = base; iexp; iexp >>= 1) {
    if (iexp & 1)
      res = res * dbl % mod;
    dbl = dbl * dbl % mod;
  }
  return res;
}

class modfactorial {
  std::vector<intmax_t> fact, fact_inv;
  intmax_t mod;

public:
  modfactorial(intmax_t N, intmax_t mod) : mod(mod) {
    fact.resize(N + 1);
    fact_inv.resize(N + 1);
    fact[0] = 1;
    for (intmax_t i = 1; i <= N; ++i)
      fact[i] = fact[i - 1] * i % mod;

    fact_inv[N] = modinv(fact[N], mod);
    for (intmax_t i = N; i--;)
      fact_inv[i] = fact_inv[i + 1] * (i + 1) % mod;
  }

  intmax_t operator()(intmax_t k) const { return fact[k]; }

  intmax_t inverse(intmax_t k) const { return fact_inv[k]; }
};

template <class Tp> std::vector<Tp> make_vector(size_t n, const Tp &x) {
  return std::vector<Tp>(n, x);
}

template <class... Ts> auto make_vector(size_t n, size_t m, const Ts &...x) {
  return std::vector<decltype(make_vector(m, x...))>(n, make_vector(m, x...));
}

std::vector<std::vector<intmax_t>> make_dp(std::vector<size_t> dig, size_t d) {
  --dig[d];
  size_t n = std::accumulate(dig.begin(), dig.end(), 0);
  size_t m = 0;
  for (size_t i = 1; i <= 4; ++i)
    m += i * dig[i];
  auto dp = make_vector(n + 1, n + 1, m + 1, intmax_t(0));
  dp[0][0][0] = 1;

  std::vector<size_t> ds;
  ds.reserve(n);
  for (size_t i = 1; i <= 4; ++i)
    for (size_t j = 0; j < dig[i]; ++j)
      ds.push_back(i);

  for (size_t i = 0; i < n; ++i) {
    for (size_t j = 0; j < n; ++j) {
      for (size_t k = 0; k <= m; ++k) {
        (dp[i + 1][j][k] += dp[i][j][k]) %= mod;
        if (k + ds[i] <= m) {
          (dp[i + 1][j + 1][k + ds[i]] += (j + 1) * dp[i][j][k]) %= mod;
        }
      }
    }
  }

  return dp[n];
}

intmax_t solve(const std::vector<intmax_t> &a) {
  size_t n = a.size();
  std::vector<size_t> dig(5);
  std::vector<intmax_t> sum(5);
  for (auto ai : a) {
    if (ai == 0) {
      ++dig[1];
      continue;
    }
    int d = 0;
    for (int i = ai; i > 0; i /= 10)
      ++d;
    ++dig[d];
    (sum[d] += ai) %= mod;
  }

  intmax_t res = 0;
  modfactorial mf(n + 1, mod);
  for (int i = 1; i <= 4; ++i) {
    if (!dig[i])
      continue;
    auto dp = make_dp(dig, i);

    for (size_t j = 0; j < n; ++j) {
      for (size_t l = 0; l + j < n; ++l) {
        intmax_t co = mf(n - 1 - j);
        (co *= mf.inverse(n - 1 - j - l)) %= mod;

        for (size_t k = 0; k < dp[j].size(); ++k) {
          (res += modmul({co, sum[i], modpow(10L, k, mod), dp[j][k]}, mod)) %=
              mod;
        }
      }
    }
  }
  return res;
}

int main() {
  size_t n;
  scanf("%zu", &n);

  std::vector<intmax_t> a(n);
  for (auto &ai : a)
    scanf("%jd", &ai);

  std::sort(a.rbegin(), a.rend());
  intmax_t res = solve(a);
  if (a[n - 1] == 0) {
    a.pop_back();
    res -= solve(a);
    if (res < 0)
      res += mod;
  }

  printf("%jd\\n", res);
}