HackerRank Combinatorics – Irresponsible Numbers

Challenge: Irresponsible Numbers

Subdomeniu: Combinatorics (combinatorics)

Scor cont: 75.0 / 75

Submission status: Accepted

Submission score: 1.0

Submission ID: 464763594

Limbaj: cpp14

Link challenge: https://www.hackerrank.com/challenges/irresponsible-numbers/problem

Cerință

Input Format

Output Format

Constraints

Cod sursă

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

static const long long MOD = 1000000007LL;

struct Mat {
    long long a[4][4];
    Mat(bool ident = false) {
        for (int i = 0; i < 4; ++i)
            for (int j = 0; j < 4; ++j)
                a[i][j] = (ident && i == j) ? 1 : 0;
    }
};

long long modpow(long long x, long long e) {
    long long r = 1 % MOD;
    x %= MOD;
    if (x < 0) x += MOD;
    while (e > 0) {
        if (e & 1LL) r = (__int128)r * x % MOD;
        x = (__int128)x * x % MOD;
        e >>= 1LL;
    }
    return r;
}

Mat mul(const Mat &A, const Mat &B) {
    Mat C;
    for (int i = 0; i < 4; ++i) {
        for (int k = 0; k < 4; ++k) if (A.a[i][k]) {
            long long aik = A.a[i][k];
            for (int j = 0; j < 4; ++j) if (B.a[k][j]) {
                C.a[i][j] = (C.a[i][j] + (__int128)aik * B.a[k][j]) % MOD;
            }
        }
    }
    return C;
}

array<long long,4> mulVec(const Mat &A, const array<long long,4> &v) {
    array<long long,4> r{0,0,0,0};
    for (int i = 0; i < 4; ++i) {
        long long s = 0;
        for (int j = 0; j < 4; ++j) {
            s = (s + (__int128)A.a[i][j] * v[j]) % MOD;
        }
        r[i] = s;
    }
    return r;
}

Mat mpow(Mat base, long long exp) {
    Mat res(true);
    while (exp > 0) {
        if (exp & 1LL) res = mul(base, res);
        base = mul(base, base);
        exp >>= 1LL;
    }
    return res;
}

Mat transDigit(int c) {
    Mat T;
    int alpha = (c <= 4) ? 1 : 0;
    int beta = (c == 9) ? 1 : 0;
    int lessA = min(c, 5);

    // state order: EA, EB, LA, LB
    // EA' = alpha*EA
    T.a[0][0] = alpha;

    // EB' = beta*EA + beta*EB
    T.a[1][0] = beta;
    T.a[1][1] = beta;

    // LA' = lessA*EA + 5*LA
    T.a[2][0] = lessA;
    T.a[2][2] = 5;

    // LB' = LA + LB
    T.a[3][2] = 1;
    T.a[3][3] = 1;

    return T;
}

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

    string xs;
    long long n;
    cin >> xs >> n;

    // Block transition for one copy of x
    Mat block(true);
    for (char ch : xs) {
        int d = ch - '0';
        Mat Td = transDigit(d);
        block = mul(Td, block);
    }

    Mat all = mpow(block, n);

    array<long long,4> v0{1,0,0,0}; // equal + in prefix-state
    auto vf = mulVec(all, v0);

    long long responsible_0_to_S = (vf[0] + vf[1] + vf[2] + vf[3]) % MOD;
    long long responsible_pos = (responsible_0_to_S - 1 + MOD) % MOD; // exclude 0

    // S = x repeated n times, modulo MOD
    long long xv = 0;
    for (char ch : xs) xv = (xv * 10 + (ch - '0')) % MOD;
    long long p10L = modpow(10, (long long)xs.size());
    long long geom;
    if (p10L == 1) {
        geom = n % MOD;
    } else {
        long long num = (modpow(p10L, n) - 1 + MOD) % MOD;
        long long denInv = modpow((p10L - 1 + MOD) % MOD, MOD - 2);
        geom = (__int128)num * denInv % MOD;
    }
    long long Smod = (__int128)xv * geom % MOD;

    long long ans = (Smod - responsible_pos) % MOD;
    if (ans < 0) ans += MOD;
    cout << ans << '\n';
    return 0;
}