HackerRank Combinatorics – k-balance number

Challenge: k-balance number

Subdomeniu: Combinatorics (combinatorics)

Scor cont: 100.0 / 100

Submission status: Accepted

Submission score: 1.0

Submission ID: 464815975

Limbaj: cpp14

Link challenge: https://www.hackerrank.com/challenges/k-balance/problem

Cerinta

Your task is to calculate the sum (indicated by S) of all the k-balanced integers between [L, R]. An integer is called k-balanced when either of #1 or #2 below holds true.

1. The length of the integer <= k     
2. Sum of the first k digits (with no leading zeros) is equal to the sum of the last k digits.    

**Input format**  
L R k

**Output format**  
S % 1000000007

**Constraints**  
0 < L <= R < 10^18  
0 < k <= 18

**Sample Input**  
9 23 1

**Sample Output**  
42

**Explanation**  
9, 11 and 22 are the only 3 integers between 9 and 23 ( both included ) that are k-balanced. Hence, the answer is 9+11+22=42

Cod sursa

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

static const int MOD = 1000000007;

struct Node {
    int cnt;
    int sum;
};

static inline int addm(int a, int b) {
    a += b;
    if (a >= MOD) a -= MOD;
    return a;
}

static inline int subm(int a, int b) {
    a -= b;
    if (a < 0) a += MOD;
    return a;
}

static inline int mulm(long long a, long long b) {
    return (int)((a * b) % MOD);
}

static int solve_len_any(int len, const vector<int>* lim) {
    // All numbers with exactly len digits (first non-zero), optionally <= lim.
    array<Node,2> cur{}, nxt{};
    cur[1] = {1, 0};

    for (int pos = 0; pos < len; ++pos) {
        nxt = {};
        for (int tight = 0; tight <= 1; ++tight) {
            int cnt = cur[tight].cnt;
            int sm = cur[tight].sum;
            if (!cnt && !sm) continue;

            int md = (tight && lim) ? (*lim)[pos] : 9;
            int lo = (pos == 0 ? 1 : 0);
            for (int d = lo; d <= md; ++d) {
                int nt = (tight && lim && d == md) ? 1 : 0;
                int ncnt = addm(nxt[nt].cnt, cnt);
                int nsum = nxt[nt].sum;
                nsum = addm(nsum, mulm(sm, 10));
                nsum = addm(nsum, mulm(cnt, d));
                nxt[nt] = {ncnt, nsum};
            }
        }
        cur = nxt;
    }

    int ans = 0;
    if (lim) ans = addm(cur[0].sum, cur[1].sum);
    else ans = cur[0].sum;
    return ans;
}

static int solve_len_balanced(int len, int k, const vector<int>* lim) {
    // len > k. Need sum(first k digits) == sum(last k digits).
    int maxAbs = 9 * k;
    int W = 2 * maxAbs + 1;
    int OFF = maxAbs;

    vector<array<Node,2>> cur(W), nxt(W);
    cur[OFF][1] = {1, 0};

    auto coeff = [&](int pos) {
        int c = 0;
        if (pos < k) c += 1;
        if (pos >= len - k) c -= 1;
        return c;
    };

    for (int pos = 0; pos < len; ++pos) {
        for (int i = 0; i < W; ++i) nxt[i] = {};
        int cf = coeff(pos);

        for (int di = 0; di < W; ++di) {
            for (int tight = 0; tight <= 1; ++tight) {
                int cnt = cur[di][tight].cnt;
                int sm = cur[di][tight].sum;
                if (!cnt && !sm) continue;

                int md = (tight && lim) ? (*lim)[pos] : 9;
                int lo = (pos == 0 ? 1 : 0);
                for (int d = lo; d <= md; ++d) {
                    int ndi = di + cf * d;
                    if (ndi < 0 || ndi >= W) continue;
                    int nt = (tight && lim && d == md) ? 1 : 0;

                    int ncnt = addm(nxt[ndi][nt].cnt, cnt);
                    int nsum = nxt[ndi][nt].sum;
                    nsum = addm(nsum, mulm(sm, 10));
                    nsum = addm(nsum, mulm(cnt, d));
                    nxt[ndi][nt] = {ncnt, nsum};
                }
            }
        }
        cur.swap(nxt);
    }

    int ans = 0;
    if (lim) ans = addm(cur[OFF][0].sum, cur[OFF][1].sum);
    else ans = cur[OFF][0].sum;
    return ans;
}

static int solve_upto(long long X, int k) {
    if (X <= 0) return 0;
    string s = to_string(X);
    int n = (int)s.size();

    int ans = 0;
    for (int len = 1; len < n; ++len) {
        if (len <= k) ans = addm(ans, solve_len_any(len, nullptr));
        else ans = addm(ans, solve_len_balanced(len, k, nullptr));
    }

    vector<int> lim(n);
    for (int i = 0; i < n; ++i) lim[i] = s[i] - '0';
    if (n <= k) ans = addm(ans, solve_len_any(n, &lim));
    else ans = addm(ans, solve_len_balanced(n, k, &lim));

    return ans;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    long long L, R;
    int k;
    cin >> L >> R >> k;

    int ans = subm(solve_upto(R, k), solve_upto(L - 1, k));
    cout << ans << '\n';
    return 0;
}