くれなゐの雑記

例を上げて 自分で手を動かして学習できる入門記事を多めに書いています

Educational Codeforces Round 53 E. Segment Sum

問題

codeforces.com

問題概要

以下の条件を満たす[l, r]の範囲の数字の和を求める問題。(場合の数ではなく、数字そのものを足す点に注意する。)

条件: 10進数表記で、使用されている数字の種類がk 以下

方針

f(A): ある数字A以下の、条件を満たす数字の和 とすると、f(r)-f(l-1)でこの問題の答えは求められる。

こういう条件を数字の数え上げは桁DPがいい。そう古事記にも書かれている。というわけで桁DPで方針を立てる。

そもそも、数え上げDPは条件を満たすDAGのパスの総数を求める問題とも言いかえることができる(参考)。 桁DPの解説は個人的にこのページが分かりやすかった。

pekempey.hatenablog.com

ということで、上記のページの要領でグラフを構築し、各辺に対し、その数字の重み(右から3桁目の4なら400の重みがつく)を付け、条件に沿うDAGのedgeの合計値を求めれば良い。

状態としては、桁数、使用した数字(10桁のbitで管理)、A未満であることが確定しているかどうか の3次元。 10桁の数字のうち、どれを使用したかをメモするためには1024必要なので、合計で18*1024*2で間に合う。

最後の桁数に関して、使用した数字の数がk以下のもののみを数え上げ、出力する。

数え上げならば、これだけで良いのだが、今回は数字の合計を求めなければならないのでもうひと工夫が必要。 ちゃんと各数字ごとに使用される回数をかけてやらなければならない。 上流側は、数字を流せば自動的に分裂して数え上げしてくれるが、下流側は分裂する回数が足りないのでキチンと数え上げをしてくれない。 例えば、一番右のケタの1は10回しか呼ばれないので雑に計算すると桁数にかかわらず1*10で10になってしまう。 でも実際には、1000までの数字の和だったら100回 1 が呼ばれる。 下流側では、上流側の数え上げた分を反映しなければならない。

なので、DPはシンプルに数え上げるだけのものと、数字そのものの合計を求めるものの2つを用意して、下記ソースコードのように遷移した。

コーナーケースとして、最上位の桁に対して0を選択し続けた場合に関して、それは0を選択していないことになる。 (000023は23扱いなので、0は含まれていない)ので、注意する。

int dp[22][(1<<10)][2];
int sums[22][(1<<10)][2];
int solve(const string str, const int K) {
    memset(dp, 0, sizeof(dp));
    memset(sums, 0, sizeof(dp));
    vector<int> nums(str.size());
    REP(i, str.size()) {
        nums[i] = str[i] - '0';
    }

    vector<int> weights(nums.size());
    weights[0] = 1;
    FOR(i, 1, nums.size()) {
        weights[i] = (weights[i - 1] * 10LL)%MOD;
    }
    reverse(ALL(weights));

    dp[0][0][0] = 1;
    sums[0][0][0] = 0;
    REP(i, nums.size()){
        int weight = weights[i];
        REP(j, (1<<10)) {
            REP(k, 2) {
                int lim = k ? 9 : nums[i];
                REP(d, lim + 1) {
                    int jj = j;
                    if (j == 1) {
                        jj = 0;
                    }
                    (dp[i + 1][jj | (1 << d)][k || d < lim] += (dp[i][j][k])) %= MOD;
                    (sums[i + 1][jj | (1<<d)][k || d < lim] += (sums[i][j][k] + (d * (weight * dp[i][j][k])%MOD)%MOD)%MOD) %= MOD;
                }
            }
        }
    }

    int res = 0;
    REP(j, (1 << 10)) REP(k, 2) {
        bitset<10> bits = j;
        if (bits.count() <= K) {
            (res += sums[nums.size()][j][k])%=MOD;
        }
    }
    return res;
}

void solve() {
    int l, r, k; cin >> l >> r >> k;
    cout << ((solve(to_string(r), k) - solve(to_string(l-1), k)) + MOD)%MOD << endl;
}