Algorithm to find Lucky Numbers

Question

I came across this question.A number is called lucky if the sum of its digits, as well as the sum of the squares of its digits is a prime number. How many numbers between A

Answer 1

You should use DP for this task. Here is my solution:

#include 

const int MAX_LENGTH = 18;
const int MAX_SUM = 162;
const int MAX_SQUARE_SUM = 1458;
int primes[1459];
long long dyn_table[19][163][1459];

void gen_primes() {
    for (int i = 0; i <= MAX_SQUARE_SUM; ++i) {
        primes[i] = 1;
    }
    primes[0] = primes[1] = 0;

    for (int i = 2; i * i <= MAX_SQUARE_SUM; ++i) {
        if (!primes[i]) {
            continue;
        }
        for (int j = 2; i * j <= MAX_SQUARE_SUM; ++j) {
            primes[i*j] = 0;
        }
    }
}

void gen_table() {
    for (int i = 0; i <= MAX_LENGTH; ++i) {
        for (int j = 0; j <= MAX_SUM; ++j) {
            for (int k = 0; k <= MAX_SQUARE_SUM; ++k) {
                dyn_table[i][j][k] = 0;
            }
        }
    }
    dyn_table[0][0][0] = 1;

    for (int i = 0; i < MAX_LENGTH; ++i) {
        for (int j = 0; j <= 9 * i; ++j) {
            for (int k = 0; k <= 9 * 9 * i; ++k) {
                for (int l = 0; l < 10; ++l) {
                    dyn_table[i + 1][j + l][k + l*l] += dyn_table[i][j][k];
                }
            }
        }
    }
}

long long count_lucky (long long max) {
            long long result = 0;
    int len = 0;
    int split_max[MAX_LENGTH];
    while (max) {
        split_max[len] = max % 10;
        max /= 10;
        ++len;
    }
    int sum = 0;
    int sq_sum = 0;
    for (int i = len-1; i >= 0; --i) {
        long long step_result = 0;
        for (int l = 0; l < split_max[i]; ++l) {
            for (int j = 0; j <= 9 * i; ++j) {
                for (int k = 0; k <= 9 * 9 * i; ++k) {
                    if (primes[j + l + sum] && primes[k + l*l + sq_sum]) {
                        step_result += dyn_table[i][j][k];
                    }
                }
            }
        }
        result += step_result;

        sum += split_max[i];
        sq_sum += split_max[i] * split_max[i];
    }

    if (primes[sum] && primes[sq_sum]) {
        ++result;
    }

    return result;
}

int main(int argc, char** argv) {
    gen_primes();
    gen_table();

    int cases = 0;
    scanf("%d", &cases);
    for (int i = 0; i < cases; ++i) {
        long long a, b;
        scanf("%lld %lld", &a, &b);
        printf("%lld\n", count_lucky(b) - count_lucky(a-1));
    }
    return 0;
}

Brief explanation:

• I'm calculating all primes up to 9 * 9 * MAX_LENGTH using Eratosthenes method;
• Later, using DP, I'm building table dyn_table where value X in dyn_table[i][j][k] means that we have exactly X numbers of length i with sum of digits equal to j and sum of its squares equal to k
• Then we can easily count amount of lucky numbers from 1 to 999..999(len times of 9). For this we just sum up all dyn_table[len][j][k] where both j and k are primes.
• To calculate amount of lucky number from 1 to random X we split interval from 1 to X into intervals with length equal to 10^K (see *count_lucky* function).
• And our last step is subtract count_lucky(a-1) (cause we are including a in our interval) from count_lucky(b).

That's all. Precalculation work for O(log(MAX_NUMBER)^3), each step have also this complexity.

I've tested my solution against linear straightforward one and results were equal