primes.cc

#include <iostream>
#include <cmath>
#include <vector>
#include <set>
#define NO_SQUARES {cout<<"NONE"<<endl;return 0;}

using namespace std;

// Enough primes.
int manyprimes[63] = {5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311, 313};

typedef vector<unsigned char> psquare;

int main() {
  int n, init;
  cin >> n >> init;

  // Sum of digits cannot be multiple of 2 or 3 for obvious reasons.
  // Sum of digits cannot be so great that it is impossible.
  // 11111 is not prime, any smaller sum is inadmissible.
  if (n % 2 == 0 || n % 3 == 0 || n - init > 36 || n < 6) NO_SQUARES;

  // Generate a set of prime numbers of the right digit sum. Trial
  // division by all odd numbers 5+ to ensure primality.
  set<int> primes;

  // Note that the sum of digits is constant only when the modulo base 9
  // is constant. So we can add 18 to the first odd prime in order to speed
  // up the prime generation.
  int startat = 10000 + ((n-1) % 9);
  if (startat % 2 == 0) {
    startat += 9;
  }
  for (int i = startat; i < 100000; i += 18) {
    if (i % 10 + (i / 10) % 10 + (i / 100) % 10 + (i / 1000) % 10 + (i / 10000) != n) continue;

    for (int p = 0; p < 63; p++) {
      if (i % manyprimes[p] == 0) goto skip;
    }

    primes.insert(i);
    skip: ;
  }

  // Make some prime caches for one- and two-digit prefixes.
  vector<vector<int> > primes_start_cache(10);
  vector<vector<int> > primes_start_cache2(100);

  // The initial-digit prime cache has the additional constraint
  // that no zeroes are allowed.
  vector<int> initprimes;

  // Fill the caches described above.
  for (set<int>::iterator it = primes.begin(); it != primes.end(); it++) {
    int d1 = *it / 10000;
    primes_start_cache[d1].push_back(*it);
    primes_start_cache2[*it / 1000].push_back(*it);
    if (d1 == init && (*it/1000) % 10 && (*it/100) % 10 && (*it/10) % 10) {
      initprimes.push_back(*it);
    }
  }

  // Make a set to store guesses in.
  set<psquare> squares;

  // Now start the guessing! O(?).
  // 1. Fill in first row with some prime, like such:
  // KPPPP
  // ?????
  // ?????
  // ?????
  // ?????
  for (int ii = 0; ii < initprimes.size(); ii++) {
    int i = initprimes[ii];
    int i2 = (i/1000) % 10;
    int i3 = (i/100) % 10;
    int i4 = (i/10) % 10;
    int i5 = i % 10;

    // 2. Fill in the first column with some prime, like such:
    // KPPPP
    // P????
    // P????
    // P????
    // P????
    for (int vv = 0; vv < initprimes.size(); vv++) {
      int v = initprimes[vv];
      int j1 = (v/1000) % 10;
      int k1 = (v/100) % 10;
      int l1 = (v/10) % 10;
      int m1 = v % 10;

      // 3. Fill in the second column with some prime, like such:
      // KPPPP
      // PP???
      // PP???
      // PP???
      // PP???
      for (int ww = 0; ww < primes_start_cache[i2].size(); ww++) {
        int w = primes_start_cache[i2][ww];
        int j2 = (w/1000) % 10;
        int k2 = (w/100) % 10;
        int l2 = (w/10) % 10;
        int m2 = w % 10;

        // 4. Fill in the second row with some prime, like such:
        // KPPPP
        // PPPPP
        // PPN?? (we can fill this in thanks to D2)
        // PP???
        // PP???
        for (int jj = 0; jj < primes_start_cache2[j1*10+j2].size(); jj++) {
          int j = primes_start_cache2[j1*10+j2][jj];
          int j3 = (j/100) % 10;
          int j4 = (j/10) % 10;
          int j5 = j % 10;

          // k3 is now known, so D2 can now be checked for primality.
          // By now we have ensured primality for: I, J, V, W, D2.
          // And of course we also ensured sum of digits for those.
          int k3 = n - i5 - j4 - l2 - m1;
          int d2 = i5 + j4 * 10 + k3 * 100 + l2 * 1000 + m1 * 10000;
          if (primes.find(d2) == primes.end()) continue;

          // Instead of going through all possible primes for m, do a brief
          // check that i3, j3, k3, i1, j2, k3, and k1, k2, k3 are capable
          // of beginning primes.
          int d1_partial = init * 10000 + j2 * 1000 + k3 * 100;
          if (primes.lower_bound(d1_partial) == primes.lower_bound(d1_partial+100)) continue;
          int xsum_partial = i3 * 10000 + j3 * 1000 + k3 * 100;
          if (primes.lower_bound(xsum_partial) == primes.lower_bound(xsum_partial+100)) continue;
          int ksum_partial = k1 * 10000 + k2 * 1000 + k3 * 100;
          if (primes.lower_bound(ksum_partial) == primes.lower_bound(ksum_partial+100)) continue;

          // 5. Fill in the last row. Using this we can fill all.
          // KPPPP
          // PPPPP
          // PPK45
          // PP123
          // PPPPP
          for (int mm = 0; mm < primes_start_cache2[m1*10+m2].size(); mm++) {
            int m = primes_start_cache2[m1*10+m2][mm];
            int m3 = (m/100) % 10;
            int m4 = (m/10) % 10;
            int m5 = m % 10;

            // Complete 4th row.
            int l3 = n - i3 - j3 - k3 - m3;
            int l4 = n - init - j2 - k3 - m5;
            int l5 = n - l1 - l2 - l3 - l4;

            // Complete 3rd row.
            int k4 = n - i4 - j4 - l4 - m4;
            int k5 = n - i5 - j5 - l5 - m5;

            // Verify K. We have already verified sum of digits for I, J, L, M, V, W, X, Y, Z, and D1/D2.
            if (k1 + k2 + k3 + k4 + k5 != n) continue;
            
            // Check primality of K, L; then X, Y, Z; then D1. (D2 was checked before).
            if (primes.find(k1*10000+k2*1000+k3*100+k4*10+k5) == primes.end()) continue;
            if (primes.find(l1*10000+l2*1000+l3*100+l4*10+l5) == primes.end()) continue;
            if (primes.find(i3*10000+j3*1000+k3*100+l3*10+m3) == primes.end()) continue;
            if (primes.find(i4*10000+j4*1000+k4*100+l4*10+m4) == primes.end()) continue;
            if (primes.find(i5*10000+j5*1000+k5*100+l5*10+m5) == primes.end()) continue;
            if (primes.find(init*10000+j2*1000+k3*100+l4*10+m5) == primes.end()) continue;

            // By now our result should be correct.
            psquare g(25);
            g[0] = init;
            g[1] = i2;
            g[2] = i3;
            g[3] = i4;
            g[4] = i5;
            g[5] = j1;
            g[6] = j2;
            g[7] = j3;
            g[8] = j4;
            g[9] = j5;
            g[10] = k1;
            g[11] = k2;
            g[12] = k3;
            g[13] = k4;
            g[14] = k5;
            g[15] = l1;
            g[16] = l2;
            g[17] = l3;
            g[18] = l4;
            g[19] = l5;
            g[20] = m1;
            g[21] = m2;
            g[22] = m3;
            g[23] = m4;
            g[24] = m5;

            // Record guess (correct!) in table.
            squares.insert(g);
          }
        }
      }
    }
  }

  if (squares.empty()) NO_SQUARES;
  
  for (set<psquare>::iterator it = squares.begin(); it != squares.end(); it++) {
    psquare g = *it;
    cout << int(g[0]) << int(g[1]) << int(g[2]) << int(g[3]) << int(g[4]) << endl;
    cout << int(g[5]) << int(g[6]) << int(g[7]) << int(g[8]) << int(g[9]) << endl;
    cout << int(g[10]) << int(g[11]) << int(g[12]) << int(g[13]) << int(g[14]) << endl;
    cout << int(g[15]) << int(g[16]) << int(g[17]) << int(g[18]) << int(g[19]) << endl;
    cout << int(g[20]) << int(g[21]) << int(g[22]) << int(g[23]) << int(g[24]) << endl;
    cout << endl;
  }
}