AOJ 2879 ごちうさ数 (RUPC 2018 day1-F)

RUPC 2018 の思い出。桁 DP。

問題へのリンク

問題概要

ごちうさ数とは、10 進法表記において「5?13」を含む正の整数のことである。? は 0〜9 のいずれでもよい。

以下の正の整数のうち、ごちうさ数が何個あるかを求めよ。

制約

• 

解法

桁 DP。
一目、51?3 のうち何桁まで合ったか、を持ちたくなる。すなわち

dp[桁][51?3 の何桁まで合ったか][smaller]

でやりたくなる。しかしこのとき、鬼のようなコーナーケースがあり、

• 5 の次が 5 (1 から 0 に戻るのではなく 1 -> 1 に戻る)
• "? = 5" からの 1 (3 から 0 に戻るのではなく 3 -> 1)

特に 2 つめがやばい。いっそ思い切って、下4桁全部もつ DP が楽に実装できる:

dp[桁][下4桁][クリアしたか][smaller]

#include <iostream>
#include <cstring>
using namespace std;

int len[5] = {5, 7, 5, 7, 7};
string N;
long long dp[21][10010][3][3]; // digit, last4, clear, smaller

bool judge(int n) {
  if ( (n/100) != 51) return false;
  if (n % 10 != 3) return false;
  return true;
}

int main() {
  while (cin >> N) {
    memset(dp, 0, sizeof(dp));
    dp[0][0][0][0] = 1;
    for (int i = 0; i <= N.size(); ++i) {
      for (int j = 0; j <= 9999; ++j) {
        for (int k = 0; k < 10; ++k) {
          int nj = (j % 1000) * 10 + k;

          // smaller -> smaller
          {
            // clear -> clear
            dp[i+1][nj][1][1] += dp[i][j][1][1];

            // nonclear -> 
            if (judge(nj)) dp[i+1][nj][1][1] += dp[i][j][0][1];
            else dp[i+1][nj][0][1] += dp[i][j][0][1];
          }

          // same -> smaller
          if (k < (int)(N[i] - '0')) {
            // clear -> clear
            dp[i+1][nj][1][1] += dp[i][j][1][0];

            // nonclear ->
            if (judge(nj)) dp[i+1][nj][1][1] += dp[i][j][0][0];
            else dp[i+1][nj][0][1] += dp[i][j][0][0];
          }

          // same -> same
          if (k == (int)(N[i] - '0')) {
            // clear -> clear
            dp[i+1][nj][1][0] += dp[i][j][1][0];

            // nonclear ->
            if (judge(nj)) dp[i+1][nj][1][0] += dp[i][j][0][0];
            else dp[i+1][nj][0][0] += dp[i][j][0][0];
          }
          
        }
      }
    }

    long long res = 0;
    for (int j = 0; j <= 9999; ++j) {
      for (int smaller = 0; smaller < 2; ++smaller) {
        res += dp[N.size()][j][1][smaller];
      }
    }

    cout << res << endl;
  }
}