USACO6.4-The Primes

The Primes
IOI'94

In the square below, each row, each column and the two diagonals can be read as a five digit prime number. The rows are read from left to right. The columns are read from top to bottom. Both diagonals are read from left to right.

+---+---+---+---+---+
| 1 | 1 | 3 | 5 | 1 |
+---+---+---+---+---+
| 3 | 3 | 2 | 0 | 3 |
+---+---+---+---+---+
| 3 | 0 | 3 | 2 | 3 |
+---+---+---+---+---+
| 1 | 4 | 0 | 3 | 3 |
+---+---+---+---+---+
| 3 | 3 | 3 | 1 | 1 |
+---+---+---+---+---+ 

• The prime numbers' digits must sum to the same number.
• The digit in the top left-hand corner of the square is pre-determined (1 in the example).
• A prime number may be used more than once in the same square.
• If there are several solutions, all must be presented (sorted in numerical order as if the 25 digits were all one long number).
• A five digit prime number cannot begin with a zero (e.g., 00003 is NOT a five digit prime number).

PROGRAM NAME: prime3

INPUT FORMAT

A single line with two space-separated integers: the sum of the digits and the digit in the upper left hand corner of the square.

SAMPLE INPUT (file prime3.in)

11 1

OUTPUT FORMAT

Five lines of five characters each for each solution found, where each line in turn consists of a five digit prime number. Print a blank line between solutions. If there are no prime squares for the input data, output a single line containing "NONE".

SAMPLE OUTPUT (file prime3.out)

The above example has 3 solutions.

11351
14033
30323
53201
13313

11351
33203
30323
14033
33311

13313
13043
32303
50231
13331

---

题目大意：给出一个5*5的正方形和左上角的一个数字，要求每一行每一列每一对角线上的数字的和都等于输入的数，并且组成的五位数为质数，输出所有可能的5*5正方形。

算法：这是一道很麻烦的暴力题，朴素的枚举必定会超时，这时候需要调换枚举顺序，以及根据枚举的五位数的位置尽可能排除不可能的情况再枚举，要优化的地方还是比较多的，详见代码。

#include <iostream>
#include <memory.h>
#include <stdio.h>
#include <vector>
#include <string>
#include <algorithm>
using namespace std;
#define MAXN 100000
#define R 0
#define C 1


void output(string str)
{
    for(int i=0; i<str.length(); i++)
    {
        printf("%c",str[i]);
        if(i%5==4)
            printf("
");
    }
}

int getDigSum(int x)
{
    int sum=0;
    while(x!=0)
    {
        sum+=x%10;
        x/=10;
    }
    return sum;
}

bool allDigOdd(int x)
{
    while(x!=0)
    {
        int dig=x%10;
        if(dig%2==0)
            return false;
        x/=10;
    }
    return true;
}

bool allNoZero(int x)
{
    while(x!=0)
    {
        int dig=x%10;
        if(dig==0)
            return false;
        x/=10;
    }
    return true;
}

int isPrime[MAXN+10],primeDig[MAXN+10][5];

int a[5][5]= {0},goalSum=0,st;
vector<int> primelist[100]; // 以d[ij]表示以i开头，j结尾的字符串
vector<int> primelistMid[1000]; // 以d[ijk]表示以i开头，j为百位，k结尾的字符串
vector<int> primelistOdd[100]; // 各位上的数都为奇数
vector<int> primelistnozero[100]; // 各位上都没有0
vector<string> res;

int isprime(int type,int num)
{

    int sum=0;
    if(type==R)
    {
        for(int i=0; i<5; i++)
            sum=sum*10+a[num][i];
    }
    else
    {
        for(int i=0; i<5; i++)
            sum=sum*10+a[i][num];
    }
    return isPrime[sum];
}

void add()
{
    string str;
    for(int i=0; i<5; i++)
        for(int j=0; j<5; j++)
        {
            str+='0'+a[i][j];
        }
    res.push_back(str);
}

void initA()
{
    for(int i=0; i<5; i++)
        for(int j=0; j<5; j++)
            a[i][j]=100;
}

int main()
{
    freopen("prime3.in","r",stdin);
    freopen("prime3.out","w",stdout);

    initA();

    // 求素数表
    memset(isPrime,-1,sizeof isPrime);
    for(int i=2; i<=MAXN; i++)
        if(isPrime[i])
            for(long long j=(long long)i*i; j<=MAXN; j+=i)
                isPrime[j]=false;


    cin>>goalSum>>st;

    for(int i=10000; i<MAXN; i++)
        if(isPrime[i] && getDigSum(i)!=goalSum)
            isPrime[i]=0;

    for(int i=10000; i<MAXN; i++)
        if(isPrime[i])
        {
            int x=i;
            for(int j=0; j<5; j++)
            {
                primeDig[i][4-j]=x%10;
                x/=10;
            }
            primelist[i/10000*10+i%10].push_back(i);
            primelistMid[i/10000*100+(i%1000)/100*10+i%10].push_back(i);
        }

    // 各位上都为奇数
    for(int i=10000; i<MAXN; i++)
        if(isPrime[i] && allDigOdd(i))
            primelistOdd[i/10000*10+i%10].push_back(i);

    // 各位上都为非0
    for(int i=10000; i<MAXN; i++)
        if(isPrime[i] && allNoZero(i))
            primelistnozero[i/10000*10+i%10].push_back(i);

    bool found=false;

    a[0][0]=st;
    // 枚举 先是右下方数字
    for(a[4][4]=1; a[4][4]<=9; a[4][4]+=2)
    {
        // 枚举左下方数字
        for(a[4][0]=1; a[4][0]<=9; a[4][0]+=2)
        {
            // 枚举右上方数字
            for(a[0][4]=1; a[0][4]<=9; a[0][4]+=2)
            {
                //填充左边质数
                for(int i=0; i<primelistnozero[st*10+a[4][0]].size(); i++)
                {
                    for(int j=0; j<5; j++)
                        a[j][0]=primeDig[primelistnozero[st*10+a[4][0]][i]][j];

                    // 填充上边质数
                    for(int ii=0; ii<primelistnozero[st*10+a[0][4]].size(); ii++)
                    {
                        for(int j=0; j<5; j++)
                            a[0][j]=primeDig[primelistnozero[st*10+a[0][4]][ii]][j];


                        // 填充下边质数
                        for(int iii=0; iii<primelistOdd[a[4][0]*10+a[4][4]].size(); iii++)
                        {
                            for(int j=0; j<5; j++)
                                a[4][j]=primeDig[primelistOdd[a[4][0]*10+a[4][4]][iii]][j];


                            //填充右边质数
                            for(int iiii=0; iiii<primelistOdd[a[0][4]*10+a[4][4]].size(); iiii++)
                            {
                                for(int j=0; j<5; j++)
                                    a[j][4]=primeDig[primelistOdd[a[0][4]*10+a[4][4]][iiii]][j];


                                // 填充中间一点
                                for(a[2][2]=0; a[2][2]<=9; a[2][2]++)
                                {

                                    // 填充主对角线
                                    for(int k=0; k<primelistMid[st*100+a[2][2]*10+a[4][4]].size(); k++)
                                    {
                                        a[1][1]=primeDig[primelistMid[st*100+a[2][2]*10+a[4][4]][k]][1];
                                        a[3][3]=primeDig[primelistMid[st*100+a[2][2]*10+a[4][4]][k]][3];

                                        // 填充辅对角线
                                        for(int l=0; l<primelistMid[a[4][0]*100+a[2][2]*10+a[0][4]].size(); l++)
                                        {
                                            a[3][1]=primeDig[primelistMid[a[4][0]*100+a[2][2]*10+a[0][4]][l]][1];
                                            a[1][3]=primeDig[primelistMid[a[4][0]*100+a[2][2]*10+a[0][4]][l]][3];

                                            int a12=goalSum-(a[1][0]+a[1][1]+a[1][3]+a[1][4]);
                                            int a21=goalSum-(a[0][1]+a[1][1]+a[3][1]+a[4][1]);
                                            int a23=goalSum-(a[0][3]+a[1][3]+a[3][3]+a[4][3]);
                                            int a32=goalSum-(a[3][0]+a[3][1]+a[3][3]+a[3][4]);
                                            a[1][2]=a12;
                                            a[2][1]=a21;
                                            a[2][3]=a23;
                                            a[3][2]=a32;
                                            if(a12>=0 && a12<10 && a21>=0 && a21<10 && a23>=0 && a23<10 && a32>=0 && a32<10
                                                    && isprime(R,1) && isprime(R,2) && isprime(R,3)
                                                    && isprime(C,1) && isprime(C,2) && isprime(C,3))
                                            {
                                                found=true;
                                                add();
                                            }
                                        }

                                    }
                                }

                            }
                        }
                    }
                }
            }
        }
    }

    sort(res.begin(),res.end());
    bool first=true;
    for(int i=0; i<res.size(); i++)
    {
        if(!first)
        {
            cout<<endl;
        }
        first=false;
        output(res[i]);
    }

    if(!found)
        printf("NONE
");
    return 0;
}