NEFU 641 素数树

素数树

Time Limit 1000ms Memory Limit 65536K

description

Suppose there is a tree named A. All nodes of A have a weight v(0 < v < 4000000).Now, we will give the definition of "Prime Node".A node is a Prime Node if the following conditions are satisfied.The subtree of A whose root node is b will be marked as B. If all nodes in B have prime weights and b's weight is greater than other nodes', then b is a Prime Node.Now you are to calculate how many Prime Nodes are in A.

input

For each test case:The fist line will contains an integer n(1 <= n <= 10000) indicating the number of nodes in A,the root of A will always be node 1.The second line has n integers giving the weights of each node numbered from 1 to n.The following n-1 lines, each contains a pair of integers x and y indicating there is an edge between x and y, give the n-1 edges of A. Process to the end of file.

output

For each test case, print the number of Prime Nodes on a single line.Print a blank line after each test case, even after the last one.

sample_input

5 12 11 5 7 3 1 2 2 3 2 5 3 4

sample_output

3

我觉得我的码可读性更强orz

#include <iostream>
#include <cstdio>
#include <cstring>

using namespace std;


const int OO=1e9;
const int INF=1e9;
const int maxn=51111;
const int max_prime=4000009;

struct NODE{
    int num;
    int max_num;
    bool visit;
    bool prime;
}tree[maxn];

struct EDGE{
    int to;
    int next;
}edges[2*maxn];

int head[maxn];
int n,edge;
int ans;

void init()
{
    memset(head,-1,sizeof(head));
    memset(edges,0,sizeof(edges));
    memset(tree,0,sizeof(tree));
    edge=0;
    ans=0;
}

void addedge(int u,int v)
{
    edges[edge].to=v;edges[edge].next=head[u];head[u]=edge++;
    edges[edge].to=u;edges[edge].next=head[v];head[v]=edge++;
}

bool not_prime[max_prime];
void find_prime()
{
    memset(not_prime,0,sizeof(not_prime));
    not_prime[0]=true;
    not_prime[1]=true;
    for (int i=2;i*i<max_prime;i++)
    {
        if (!not_prime[i])
        {
            for (int j=i+i;j<max_prime;j+=i)
            {
                not_prime[j]=true;
            }
        }
    }
}

void dfs(int u)
{
    int v;
    bool ok=tree[u].prime;
    tree[u].visit=true;
    for (int i=head[u];i!=-1;i=edges[i].next)
    {
        v=edges[i].to;
        if (!tree[v].visit)
        {
            dfs(v);
            if (!tree[v].prime)
            {
                tree[u].prime=false;
                ok=false;
            }
            if (tree[v].max_num>=tree[u].num)
            {
                ok=false;
            }
            if (tree[v].max_num>=tree[u].max_num)
            {
                tree[u].max_num=tree[v].max_num;
            }
        }
    }
    if (ok) ans++;
}

int main()
{
    find_prime();
    while (scanf("%d",&n)!=EOF)
    {
        init();
        for (int i=1;i<=n;i++)
        {
            scanf("%d",&tree[i].num);
            if ( !not_prime[tree[i].num] )
            {
                tree[i].prime=true;
            }
            else
            {
                tree[i].prime=false;
            }
            tree[i].visit=false;
            tree[i].max_num=tree[i].num;
        }
        for (int i=1;i<n;i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            addedge(u,v);
        }
        dfs(1);
        printf("%d\n\n",ans);
    }
    return 0;
}