けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 154 E - Almost Everywhere Zero (500 点)

DP しなくても間に合うけど、桁 DP 的な考え方が役に立つ!

問題へのリンク

問題概要

100 桁以下の整数  S が与えられる。 1 以上  S 以下の整数であって、十進法表記で 0 以外の数値がちょうど  K 個であるようなものが何個あるのかを求めよ。

制約

  •  S の桁数  \le 100
  •  1 \le K \le 3

考えたこと

一般に大きな整数  S が与えられて、それ以下の整数で条件を満たすもの全体について調べたいとき、こんな感じの考え方が使える!!!

drken1215.hatenablog.com

たとえば、S = 314159 のときは、まずは

  • 000000 〜 099999
  • 100000 〜 299999
  • 300000 〜 314159

とに場合分けして考えることにする。そこで再帰的に解くために


rec(i, k, smaller) := i 桁目以降について、0 以外の値を残り K 個使用可能という状況を考える。このとき i 桁目までの部分が「等しい」か「strict に小さくなっている」かを smaller フラグによって分岐する。


という関数を用意してあげる。答えは rec(0, K, false) となる。上記の 3 つの場合についてはそれぞれ

  • rec(1, K, true) 通り
  • rec(1, K-1, true) × (S[0] - 1) 通り
  • rec(1, K-1, false) 通り

という風になる。また smaller = true になった瞬間、答えは求められる。S の桁数を  N として、

  • rec(i, k, true) = C(N-i, k) × 9k

となる。

以上の再帰関数を実装すれば求められる。メモ化すると桁 DP になる。今回はメモ化しなくても大丈夫。

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

long long com(long long N, long long R) {
    if (R < 0 || R > N) return 0;
    if (R == 1) return N;
    else if (R == 2) return N * (N-1) / 2;
    else return N * (N-1) * (N-2) / 6;
}

long long pow(long long N, long long k) {
    long long res = 1;
    for (int i = 0; i < k; ++i) res *= N;
    return res;
}

string S;
int N, K;

long long solve(int i, int k, int smaller) {
    if (i == N) {
        if (k == 0) return 1;
        else return 0;
    }
    if (k == 0) return 1;
    
    if (smaller) return com(N-i, k) * pow(9, k);
    else {
        if (S[i] == '0') return solve(i+1, k, false);
        else {
            long long zero = solve(i+1, k, true);
            long long aida = solve(i+1, k-1, true) * (S[i] - '1');
            long long icchi = solve(i+1, k-1, false);
            return zero + aida + icchi;
        }
    }
}

int main() {
    cin >> S >> K;
    N = S.size();
    cout << solve(0, K, false) << endl;
}