HackerRank Lucky Numbers problem solution

In this HackerRank Lucky Numbers problem solution, 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 and B inclusive, are lucky?

Problem solution in Java.

import java.io.*;
import java.math.*;
import java.text.*;
import java.util.*;
import java.util.regex.*;

public class Solution {

    /*
     * Complete the luckyNumbers function below.
     */
    static long luckyNumbers(long a, long b) {
        return Method1.luckyNumbers(a, b);
    }

    static class Method1 {
        static final int max_bit = 19;
        static final int max_digit_sum = 9 * 18 + 1;
        static final int max_squre_digit_sum = 9 * 9 * 18 + 1;

        static long[][][] sum = new long[max_bit][max_digit_sum][max_squre_digit_sum];
        static long[][][] left_sum = new long[max_bit][max_digit_sum][max_squre_digit_sum];
        static boolean[] is_prime = new boolean[max_squre_digit_sum];
        static long[] bit_sum = new long[max_bit];
        static int tot = 231, max1, max2;
        static int[] prime = new int[tot];
        static {
            Arrays.fill(is_prime, true);
            int ct = 0;
            is_prime[0] = is_prime[1] = false;
            for (int i = 2; i < max_squre_digit_sum; i++)
                if (is_prime[i]) {
                    for (int j = i + i; j < max_squre_digit_sum; j += i)
                        is_prime[j] = false;
                    prime[ct++] = i;
                    if (i < max_digit_sum)
                        max1 = Math.max(max1, i);
                    max2 = Math.max(max2, i);
                }
            sum[0][0][0] = 1;
            for (int i = 0; prime[i] <= max1; i++)
                for (int j = 0; j < tot && prime[j] <= max2; j++) {
                    left_sum[0][prime[i]][prime[j]] += 1;
                }

            for (int i = 0; i < max_bit - 1; i++) {
                for (int next = 0; next < 10; next++) {

                    for (int j = 0; j + next < max_digit_sum; j++)
                        for (int k = 0; k + next * next < max_squre_digit_sum; k++) {
                            sum[i + 1][j + next][k + next * next] += sum[i][j][k];
                            if (next > 0 && is_prime[j + next] && is_prime[k + next * next])
                                bit_sum[i + 1] += sum[i][j][k];
                        }

                    for (int j = next; j < max_digit_sum; j++)
                        for (int k = next * next; k < max_squre_digit_sum; k++) {
                            left_sum[i + 1][j - next][k - next * next] += left_sum[i][j][k];
                        }
                }
            }
        }

        static long luckyNumbers(long a, long b) {
            return go(b) - go(a - 1);
        }

        static long go(long N) {
            if (N == 0)
                return 0;
            long ret = 0;
            boolean first = true;
            int pre_digit_sum = 0, pre_sdigit_sum = 0;
            for (int i = 19; i > 0; i--) {
                int bit = get_bit(N, i - 1);
                int least;
                if (bit != 0 && first) {
                    least = 1;
                    first = false;
                    for (int j = 1; j < i; j++)
                        ret += bit_sum[j];
                } else {
                    least = 0;
                }

                for (int nbit = least; nbit < bit; nbit++) {
                    int digit_sum = pre_digit_sum + nbit;
                    int sdigit_sum = pre_sdigit_sum + nbit * nbit;
                    ret += left_sum[i - 1][digit_sum][sdigit_sum];
                }
                pre_digit_sum += bit;
                pre_sdigit_sum += bit * bit;
            }
            if (is_prime[pre_digit_sum] && is_prime[pre_sdigit_sum])
                ret += 1;
            return ret;
        }

        static int get_bit(long n, int m) {
            for (int i = 0; i < m; i++)
                n /= 10;
            return (int) (n % 10);
        }
    }
    private static final Scanner scanner = new Scanner(System.in);

    public static void main(String[] args) throws IOException {
        BufferedWriter bufferedWriter = new BufferedWriter(new FileWriter(System.getenv("OUTPUT_PATH")));

        int t = Integer.parseInt(scanner.nextLine().trim());

        for (int tItr = 0; tItr < t; tItr++) {
            String[] ab = scanner.nextLine().split(" ");

            long a = Long.parseLong(ab[0].trim());

            long b = Long.parseLong(ab[1].trim());

            long result = luckyNumbers(a, b);

            bufferedWriter.write(String.valueOf(result));
            bufferedWriter.newLine();
        }

        bufferedWriter.close();
    }
}

Problem solution in C++.

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<string>
#include<vector>
#include<map>
#include<list>
#include<queue>
#include<set>
using namespace std;
typedef vector<int> VI;
typedef long long LL;
#define FOR(x, b, e) for(int x=b; x<=(e); ++x)
#define FORD(x, b, e) for(int x=b; x>=(e); --x)
#define REP(x, n) for(int x=0; x<(n); ++x)
#define VAR(v,n) __typeof(n) v=(n)
#define ALL(c) c.begin(),c.end()
#define SIZE(x) (int)x.size()
#define FOREACH(i,c) for(VAR(i,(c).begin());i!=(c).end();++i)
#define PB push_back
#define ST first
#define ND second
const int MILION=1000*1000;
const int MAX_SUMA=200;
const int MAX_KWADRAT=2000;
const int MAX_CYFR=20;

LL pot[MAX_CYFR];
bool pierwsza[MAX_SUMA][MAX_KWADRAT];
LL w[MAX_SUMA][MAX_KWADRAT][MAX_CYFR];


bool czyPierwsza(int k){
    if (k<2) return false;
    int i=2;
    while (i*i<=k){
        if (k%i==0) return false;
        i++;
    }
    return true;
}

void ustalPot(){
    pot[0]=1;
    FOR(i,1,MAX_CYFR-1){
        pot[i]=pot[i-1]*10;
    }
}

void ustalPierwsze(){
    FOR(i,2,MAX_SUMA-2){
        FOR(j,2,MAX_KWADRAT-2){
            if (czyPierwsza(i) && czyPierwsza(j)){
                pierwsza[i][j]=true;
            }
        }
    }
}

void ustalW(){
    REP(i,MAX_SUMA){
        REP(j,MAX_KWADRAT){
            if (pierwsza[i][j]){
                w[i][j][0]++;
            }
        }
    }
    FOR(p,1,MAX_CYFR-1){
        REP(i,MAX_SUMA-2){
            REP(j,MAX_KWADRAT-2){
                REP(k,10){
                    if (i+k<MAX_SUMA-2 && j+k*k < MAX_KWADRAT-2){
                        w[i][j][p]+=w[i+k][j+k*k][p-1];
                    }
                }
            }
        }
    }
}

LL daj(int c, int k, LL T){
    if (T<1){
        return pierwsza[c][k];
    }
    LL wyn=0;
    int wyk=0;
    int cyf=1;
    while (pot[wyk]<=T) wyk++;
    wyk--;
    while (cyf*pot[wyk]<=T) cyf++;
    cyf--;
    wyn+=daj(c+cyf,k+cyf*cyf,T-cyf*pot[wyk]);
    if (wyk>-1){
        FOR(i,0,cyf-1){
            wyn+=w[c+i][k+i*i][wyk];
        }
    }
    return wyn;
}


int main(){
    ustalPot();
    ustalPierwsze();
    ustalW();
    int t;
    LL a,b;
    scanf("%d",&t);
    REP(i,t){
        scanf("%lld%lld",&a,&b);
        printf("%lldn",daj(0,0,b)-daj(0,0,a-1));
    }
    return 0;
}

Problem solution in C.

/* Enter your code here. Read input from STDIN. Print output to STDOUT */
#include <stdio.h>
#include <string.h>


int isprime[10000];
unsigned long long dp[19][2*82][2*730];
unsigned long long ans[19][10][164][1460];

int start_sum_sqr[19][163];
int end_sum_sqr[19][163];

void prime(){
	memset(isprime,0,10000*4);
	isprime[0]=1;
	isprime[1]=1;	
	for (int i  = 2; i < 10000; i += 1){
		if(isprime[i]==0){
			for (int j  = i*i; j < 10000; j += i){
				isprime[j]=1;
			}	
		}
	}
}





void inc(int i,int j,int k, int val){
	if(i<19 && j<164 && k< 1460)
		dp[i][j][k] += val;
}


unsigned long long setDP(){

	memset(dp, 0, sizeof(dp[0][0][0])*19*82*730);
	
	dp[0][0][0]=1;
    for (int i = 0; i < 18; ++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) {
                    dp[i + 1][j + l][k + l*l] += dp[i][j][k];
                }
            }
        }
    }

}

unsigned long long solve(unsigned long long num){
	if(num<=0)return 0;
	int digs[20],len=0;
	while(num){
		digs[len++]=(num%10); 
		num=num/10;
	}
	
	int sum=0,sqr_sum=0;
	unsigned long long ret=0;
	for (int i  = len-1; i >= 0; i -= 1){
		int dig = digs[i];	
		
		for (int j  = 0; j < dig; j += 1){
			if(ans[i][j][sum][sqr_sum]!=0){
				ret+=ans[i][j][sum][sqr_sum];
				continue;
			}
			unsigned long long  x=0;
			for (int k  = 0; k <= 9*i; k += 1){
				if(isprime[k+sum+j]) continue;
				for (int l  = start_sum_sqr[i][k]; l <= end_sum_sqr[i][k]; l += 1){
					if(isprime[l+ j*j + sqr_sum]==0) x+=dp[i][k][l];
				}
			}
			ans[i][j][sum][sqr_sum] = x;
			ret+=x;
		}
		sum+=dig;
		sqr_sum+=(dig*dig);
	}
	if(isprime[sum]==0 && isprime[sqr_sum]==0) ret++;
	return ret;
}

void set_sqrs(){
	for (int i = 0; i < 19; i += 1){
		for (int j = 0; j < 163; j += 1){
			for (int k = 0; k < 1460; k += 1){
				if(dp[i][j][k]!=0){
					start_sum_sqr[i][j]=k;
					break;
				}
			}
			for (int k = 1459; k>=0; k -= 1){
				if(dp[i][j][k]!=0){
					end_sum_sqr[i][j]=k;
					break;
				}
			}
		}
	}
}

int main(){
	prime();
	setDP();	
	set_sqrs();
	unsigned long long tc,a,b; 
	scanf("%lld",&tc);
	while(tc--){
		scanf("%lld %lld", &a, &b);
		if(b == 1000000000000000000ll)b--;
		printf("%lldn", solve(b) - solve(a-1));
	}
	return 0;
}