square869120Contest #2 C: 何通りの分割方法がある？

問題

C: 何通りの分割方法がある？ - square869120Contest #2 | AtCoder

解法

文字列が短いので、全ての部分文字列を数字にした時の値を計算しておく。

メモ化再帰で 「substr(pos, length) で D 以下になる組み合わせの数」を求めれば良い。

コード

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1000000007;
map<tuple<int, int, int>, int> dp;
string T;
vector<vector<int>> s_atoi;

int substr_atoi(int pos, int length) {
  assert(length > 0);
  return s_atoi[pos][pos + length - 1];
}

int64_t dfs(int D, int pos, int length) {
  if (length == 1) {
    if (D >= (T[pos] - '0')) return 1;
    return 0;
  }
  if (dp.count(make_tuple(D, pos, length)))
    return dp[make_tuple(D, pos, length)];
  int64_t ans = 0;
  if (substr_atoi(pos, length) <= D) ans++;
  for (int i = 1; i < length; ++i) {
    int p = substr_atoi(pos, i);
    if (p > D) break;
    ans += dfs(D - p, pos + i, length - i);
    if (ans > MOD) ans -= MOD;
  }
  dp[make_tuple(D, pos, length)] = ans;
  return ans;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cin >> T;
  int D;
  cin >> D;
  s_atoi.assign(T.size(), vector<int>(T.size(), 0));
  for (int i = 0; i < T.size(); ++i) {
    int num = 0;
    for (int j = i; j < T.size(); ++j) {
      if (num < 0) {
        s_atoi[i][j] = D + 1;
        continue;
      }
      num = num * 10 + (T[j] - '0');
      if (num > D) num = D + 1;
      s_atoi[i][j] = num;
    }
  }
  cout << (dfs(D, 0, T.size()) % MOD) << endl;
}