usaco 4.3 the primes 2010.8.6

/*
ID:
PROG: prime3
LANG: C++
*/

/** 先求出總和是合法的質數，枚舉斜角 + 井字，剩下的空格用減法，使用binary search查找 **/

/**
Executing...
Test 1: TEST OK [0.011 secs, 2948 KB]
Test 2: TEST OK [0.011 secs, 2948 KB]
Test 3: TEST OK [0.032 secs, 2948 KB]
Test 4: TEST OK [0.065 secs, 2948 KB]
Test 5: TEST OK [0.108 secs, 2948 KB]
Test 6: TEST OK [0.151 secs, 2948 KB]
Test 7: TEST OK [0.270 secs, 2948 KB]
Test 8: TEST OK [0.464 secs, 3080 KB]
Test 9: TEST OK [0.518 secs, 2948 KB]
Test 10: TEST OK [0.724 secs, 3080 KB]

All tests OK.
**/

#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

inline bool legal_digit_sum(vector<int> &D, unsigned int value,
const int &S)
{
int multiple[5] = {10000, 1000, 100, 10, 1}, sum = 0;
for (unsigned int i = 0; i != 5; ++i)
{
D[i] = (value / multiple[i]) % 10;
sum += D[i];
}
return(sum == S);
}

void quicksort_p(vector< vector<unsigned int> > &A, int left, int right, unsigned int x)
{
if (left < right)
{
int i = right + 1, j = left;
while (true)
{
while (i > j && A[--i][x] > A[left][x])
;
while (i > j && A[++j][x] < A[left][x])
;
swap(A[i][x], A[j][x]);
if (i == j)
break;
}
swap(A[left][x], A[j][x]);
quicksort_p(A, left, j - 1, x);
quicksort_p(A, j + 1, right, x);
}
}

inline void generate_primes(vector< vector<int> > &P,
vector<unsigned int> &PI,
vector< vector<unsigned int> > &PP,
const int &S)
{
const unsigned int SIZE = 1e5 + 1;
vector<bool> num(SIZE, true);
for (unsigned int i = 4; i < SIZE; i += 2)
num[i] = false;
for (unsigned int i = 6; i < SIZE; i += 3)
num[i] = false;
num[0] = num[1] = false;
for (unsigned int i = 5; i < SIZE; i += 4)
{
if (num[i])
{
if ((unsigned int long long)i * i < SIZE)
for (unsigned int j = (i << 1), k = i * i; k < SIZE; k += j)
if (num[k])
num[k] = false;
}
i += 2;
if (i < SIZE && num[i])
{
if ((unsigned int long long)i * i < SIZE)
for (unsigned int j = (i << 1), k = i * i; k < SIZE; k += j)
if (num[k])
num[k] = false;
}
}
vector<int> digit(5);
vector<unsigned int> partial(7);
/**
* partial index
* 0 - index04
* 1 - index13
* 2 - index2
* 3 - index014
* 4 - index12
* 5 - index0134
* 6 - index123
* **/
int multiple[5] = {10000, 1000, 100, 10, 1};
for (unsigned int i = 10000; i != SIZE; ++i)
if (num[i] && legal_digit_sum(digit, i, S))
{
P.push_back(digit);
PI.push_back(i);
partial[0] = digit[0] * multiple[0] + digit[4] * multiple[4];
partial[1] = digit[1] * multiple[1] + digit[3] * multiple[3];
partial[2] = digit[2] * multiple[2];
partial[3] = digit[0] * multiple[0] + digit[1] * multiple[1] + digit[4] * multiple[4];
partial[4] = digit[1] * multiple[1] + digit[2] * multiple[2];
partial[5] = digit[0] * multiple[0] + digit[1] * multiple[1] +
digit[3] * multiple[3] + digit[4] * multiple[4];
partial[6] = digit[1] * multiple[1] + digit[2] * multiple[2] + digit[3] * multiple[3];
PP.push_back(partial);
}
for (unsigned int i = 0; i != 7; ++i)
quicksort_p(PP, 0, PP.size() - 1, i);
}

bool binary_search_int(const vector<unsigned int> &PI,
int low, int high, const unsigned int &key)
{
if (low <= high)
{
int mid = (low + high) >> 1;
if (PI[mid] == key)
return(true);
else if (PI[mid] < key)
return(binary_search_int(PI, mid + 1, high, key));
else
return(binary_search_int(PI, low, mid - 1, key));
}
return(false);
}

bool binary_search_p(const vector< vector<unsigned int> > &PP,
int low, int high,
const unsigned int &key, const unsigned int &x)
{
if (low <= high)
{
int mid = (low + high) >> 1;
if (PP[mid][x] == key)
return(true);
else if (PP[mid][x] < key)
return(binary_search_p(PP, mid + 1, high, key, x));
else
return(binary_search_p(PP, low, mid - 1, key, x));
}
return(false);
}

inline bool check_diagonal(const vector< vector<int> > &square,
const vector< vector<int> > &P,
const vector< vector<unsigned int> > &PP)
{
unsigned int p1, multiple[5] = {10000, 1000, 100, 10, 1};
p1 = square[0][0] * multiple[0] + square[4][0] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 0))
return(false);
p1 = square[0][4] * multiple[0] + square[4][4] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 0))
return(false);
p1 = square[0][0] * multiple[0] + square[0][4] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 0))
return(false);
p1 = square[4][0] * multiple[0] + square[4][4] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 0))
return(false);
p1 = square[1][1] * multiple[1] + square[1][3] * multiple[3];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 1))
return(false);
p1 = square[3][1] * multiple[1] + square[3][3] * multiple[3];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 1))
return(false);
p1 = square[1][1] * multiple[1] + square[3][1] * multiple[3];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 1))
return(false);
p1 = square[1][3] * multiple[1] + square[3][3] * multiple[3];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 1))
return(false);
p1 = square[2][2] * multiple[2];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 2))
return(false);
return(true);
}

inline bool check_row1(const int &sum,
const vector< vector<int> > &square,
const vector< vector<unsigned int> > &PP)
{
vector<int> p_sum(5, 0);
p_sum[0] = square[0][0] + square[1][0] + square[4][0];
p_sum[4] = square[0][4] + square[1][4] + square[4][4];
p_sum[1] = 0;
p_sum[2] = square[1][2] + square[2][2];
p_sum[3] = 0;
for (unsigned int k = 0; k != 5; ++k)
if (p_sum[k] > sum)
return(false);
unsigned int p1, multiple[5] = {10000, 1000, 100, 10, 1};
p1 = square[0][0] * multiple[0] + square[1][0] * multiple[1] + square[4][0] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 3))
return(false);
p1 = square[0][4] * multiple[0] + square[1][4] * multiple[1] + square[4][4] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 3))
return(false);
p1 = square[1][2] * multiple[1] + square[2][2] * multiple[2];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 4))
return(false);
return(true);
}

inline bool check_row3(const int &sum,
const vector< vector<int> > &square,
const vector< vector<unsigned int> > &PP)
{
vector<int> p_sum(5, 0);
p_sum[0] = square[0][0] + square[1][0] + square[3][0] + square[4][0];
p_sum[4] = square[0][4] + square[1][4] + square[3][4] + square[4][4];
p_sum[1] = 0;
p_sum[2] = square[1][2] + square[2][2] + square[3][2];
p_sum[3] = 0;
for (unsigned int k = 0; k != 5; ++k)
if (p_sum[k] > sum)
return(false);
unsigned int p1, multiple[5] = {10000, 1000, 100, 10, 1};
p1 = square[0][0] * multiple[0] + square[1][0] * multiple[1] +
square[3][0] * multiple[3] + square[4][0] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 5))
return(false);
p1 = square[0][4] * multiple[0] + square[1][4] * multiple[1] +
square[3][4] * multiple[3] + square[4][4] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 5))
return(false);
p1 = square[1][2] * multiple[1] + square[2][2] * multiple[2] +
square[3][2] * multiple[3];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 6))
return(false);
return(true);
}

inline bool check_col1(const int &sum,
const vector< vector<int> > &square,
const vector< vector<unsigned int> > &PP)
{
vector<int> p_sum(5, 0);
p_sum[0] = square[0][0] + square[0][1] + square[0][4];
p_sum[4] = square[4][0] + square[4][1] + square[4][4];
p_sum[1] = 0;
p_sum[2] = square[2][1] + square[2][2];
p_sum[3] = 0;
for (unsigned int k = 0; k != 5; ++k)
if (p_sum[k] > sum)
return(false);
unsigned int p1, multiple[5] = {10000, 1000, 100, 10, 1};
p1 = square[0][0] * multiple[0] + square[0][1] * multiple[1] + square[0][4] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 3))
return(false);
p1 = square[4][0] * multiple[0] + square[4][1] * multiple[1] + square[4][4] * multiple[4];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 3))
return(false);
p1 = square[2][1] * multiple[1] + square[2][2] * multiple[2];
if (!binary_search_p(PP, 0, PP.size() - 1, p1, 4))
return(false);
return(true);
}

inline bool check_col3(const int &sum,
vector< vector<int> > &square,
const vector< vector<int> > &P,
const vector<unsigned int> &PI)
{
vector<int> p_sum(5, 0);
p_sum[0] = square[0][0] + square[0][1] + square[0][3] + square[0][4];
p_sum[4] = square[4][0] + square[4][1] + square[4][3] + square[4][4];
p_sum[1] = 0;
p_sum[2] = square[2][1] + square[2][2] + square[2][3];
p_sum[3] = 0;
for (unsigned int k = 0; k != 5; ++k)
if (p_sum[k] > sum)
return(false);
square[0][2] = sum - (square[0][0] + square[0][1] + square[0][3] + square[0][4]);
if (square[0][2] <= 0)
return(false);
unsigned int i, p1 = 0, multiple[5] = {10000, 1000, 100, 10, 1};
for (i = 0; i != 5; ++i)
p1 += square[0][i] * multiple[i];
if (!binary_search_int(PI, 0, PI.size() - 1, p1))
return(false);
int v1, v2;
v1 = sum - (square[0][2] + square[1][2] + square[2][2] + square[3][2]);
v2 = sum - (square[4][0] + square[4][1] + square[4][3] + square[4][4]);
if (v1 != v2)
return(false);
if (v1 <= 0)
return(false);
square[4][2] = v1;
p1 = 0;
for (i = 0; i != 5; ++i)
p1 += square[4][i] * multiple[i];
if (!binary_search_int(PI, 0, PI.size() - 1, p1))
return(false);
p1 = 0;
for (i = 0; i != 5; ++i)
p1 = p1 * 10 + square[i][2];
if (!binary_search_int(PI, 0, PI.size() - 1, p1))
return(false);
square[2][0] = sum - (square[0][0] + square[1][0] + square[3][0] + square[4][0]);
if (square[2][0] <= 0)
return(false);
p1 = 0;
for (i = 0; i != 5; ++i)
p1 = p1 * 10 + square[i][0];
if (!binary_search_int(PI, 0, PI.size() - 1, p1))
return(false);
v1 = sum - (square[2][0] + square[2][1] + square[2][2] + square[2][3]);
v2 = sum - (square[0][4] + square[1][4] + square[3][4] + square[4][4]);
if (v1 != v2)
return(false);
if (v1 <= 0)
return(false);
square[2][4] = v1;
p1 = 0;
for (i = 0; i != 5; ++i)
p1 += square[2][i] * multiple[i];
if (!binary_search_int(PI, 0, PI.size() - 1, p1))
return(false);
p1 = 0;
for (i = 0; i != 5; ++i)
p1 = p1 * 10 + square[i][4];
if (!binary_search_int(PI, 0, PI.size() - 1, p1))
return(false);
return(true);
}

inline bool operator < (const vector< vector<int> > &s1,
const vector< vector<int> > &s2)
{
for (unsigned int i = 0; i != 5; ++i)
for (unsigned int j = 0; j != 5; ++j)
{
if (s1[i][j] < s2[i][j])
return(true);
else if (s1[i][j] > s2[i][j])
return(false);
}
return(false);
}

inline bool operator > (const vector< vector<int> > &s1,
const vector< vector<int> > &s2)
{
for (unsigned int i = 0; i != 5; ++i)
for (unsigned int j = 0; j != 5; ++j)
{
if (s1[i][j] > s2[i][j])
return(true);
else if (s1[i][j] < s2[i][j])
return(false);
}
return(false);
}

void quicksort_sol(vector< vector< vector<int> > > &A, int left, int right)
{
if (left < right)
{
int i = right + 1, j = left;
while (true)
{
while (i > j && A[--i] > A[left])
;
while (i > j && A[++j] < A[left])
;
A[i].swap(A[j]);
if (i == j)
break;
}
A[left].swap(A[j]);
quicksort_sol(A, left, j - 1);
quicksort_sol(A, j + 1, right);
}
}

int main()
{
ofstream fout ("prime3.out");
ifstream fin ("prime3.in");
int sum, first_digit;
fin >> sum >> first_digit;
if (sum > 45)
{
fout << "NONE\n";
return(0);
}
vector<unsigned int> primes_int;
vector< vector<int> > primes;
vector< vector<unsigned int> > primes_partial;
vector< vector< vector<int> > > solution;
primes.clear(), primes_int.clear(), solution.clear(), primes_partial.clear();
generate_primes(primes, primes_int, primes_partial, sum);
unsigned int size = primes.size();
vector<int> row_s(5, 0);
vector< vector<int> > square(5, row_s);
vector<unsigned int> r(6);
unsigned int i, start;
for (start = 0; start < size && primes[start][0] < first_digit; ++start)
;
for (r[0] = start; r[0] < size; ++r[0])  // diagonal 1
{
if (primes[r[0]][0] > first_digit)
break;
for (i = 0; i != 5; ++i)
square[i][i] = primes[r[0]][i];
for (r[1] = 0; r[1] < size; ++r[1])  // diagonal 2
{
if (!(primes[r[1]][2] == square[2][2]))
continue;
for (i = 0; i != 5; ++i)
square[4 - i][i] = primes[r[1]][i];
if (!check_diagonal(square, primes, primes_partial))
continue;
for (r[2] = 0; r[2] < size; ++r[2])  // row 1
{
if (!(primes[r[2]][1] == square[1][1] && primes[r[2]][3] == square[1][3]))
continue;
square[1] = primes[r[2]];
if (!check_row1(sum, square, primes_partial))
continue;
for (r[3] = 0; r[3] < size; ++r[3])  // row 3
{
if (!(primes[r[3]][1] == square[3][1] && primes[r[3]][3] == square[3][3]))
continue;
square[3] = primes[r[3]];
if (!check_row3(sum, square, primes_partial))
continue;
for (r[4] = 0; r[4] < size; ++r[4])  // col 1
{
if (!(primes[r[4]][1] == square[1][1] && primes[r[4]][3] == square[3][1]))
continue;
for (i = 0; i != 5; ++i)
square[i][1] = primes[r[4]][i];
if (!check_col1(sum, square, primes_partial))
continue;
for (r[5] = 0; r[5] < size; ++r[5])  // col 3
{
if (!(primes[r[5]][1] == square[1][3] && primes[r[5]][3] == square[3][3]))
continue;
for (i = 0; i != 5; ++i)
square[i][3] = primes[r[5]][i];
if (check_col3(sum, square, primes, primes_int))
solution.push_back(square);
}
}
}
}
}
}
if (!solution.size())
fout << "NONE\n";
else
{
quicksort_sol(solution, 0, solution.size() - 1);
for (unsigned int i = 0; i != solution.size(); ++i)
{
if (i)
fout << "\n";
for (unsigned int j = 0; j != 5; ++j)
{
for (unsigned int k = 0; k != 5; ++k)
fout << solution[i][j][k];
fout << "\n";
}
}
}
return(0);
}