CP Notebook

← all categories

Digit Dp a2d99e57

Digit DP skeleton: count integers in [0, x] satisfying some digit-local constraint. State is [pos, extra, tight, leadingZero]; extra is problem-specific side state (e.g. "previous digit" for a no-adjacent-equal-digits constraint, sized here as an int up to 11 — 0-9 plus a 10 = "no previous digit yet" sentinel — adjust EXTRA and the transition condition for your problem). tight means the prefix so far equals x's prefix exactly (so the next digit is capped); leadingZero means every digit so far has been a leading zero (not yet "really started"). countUpTo(r) - countUpTo(l-1) gives [l, r]; countUpTo(-1) = 0 by convention (call with l-1 possibly negative).

Time: O(digits \times EXTRA \times 2 \times 2 \times 10) states. tested

content/dp/digit-dp.h

struct DigitDP {
    static constexpr int EXTRA = 11;  // problem-specific side-state cardinality
    vector<int> d;
    ll dp[20][EXTRA][2][2];
    bool vis[20][EXTRA][2][2];

    ll calc(int pos, int extra, bool tight, bool leadingZero) {
        if (pos == (int)d.size()) return 1;
        if (vis[pos][extra][tight][leadingZero]) return dp[pos][extra][tight][leadingZero];
        vis[pos][extra][tight][leadingZero] = true;

        int ub = tight ? d[pos] : 9;
        ll res = 0;
        for (int c = 0; c <= ub; c++) {
            if (leadingZero && c == 0) {
                res += calc(pos + 1, EXTRA - 1, tight && c == d[pos], true);
            } else if (c != extra) {  // op: the actual digit-local constraint goes here
                res += calc(pos + 1, c, tight && c == d[pos], false);
            }
        }
        return dp[pos][extra][tight][leadingZero] = res;
    }

    ll countUpTo(ll x) {
        if (x < 0) return 0;
        memset(dp, 0, sizeof(dp));
        memset(vis, 0, sizeof(vis));
        d.clear();
        if (x == 0) return 1;
        while (x) { d.push_back((int)(x % 10)); x /= 10; }
        reverse(d.begin(), d.end());
        return calc(0, EXTRA - 1, true, true);
    }
};