POJ 2449 Remmarguts' Date (第k短路 A*搜索算法模板)

#include <cstdio>
#include <cstring>
#include <queue>
#include <vector>
using namespace std;
int const INF = 0xfffffff;
int const MAX1 = 1005;
int const MAX2 = 1e5 + 5;

int dis[MAX1], head[MAX1], head2[MAX1];
int n, m, cnt;
bool vis[MAX1];

struct node
{
    int v, w;
    int next;
}e[MAX2], e2[MAX2];

struct node1
{
    int v, g, f;
    bool operator < (const node1 &r) const
    {
        if(r.f == f) 
            return r.g < g;
        return r.f < f;
    }
};

void Add(int u, int v, int w)
{
    e[cnt].v = v;
    e[cnt].w = w;
    e[cnt].next = head[u];
    head[u] = cnt;

    e2[cnt].v = u;
    e2[cnt].w = w;
    e2[cnt].next = head2[v];
    head2[v] = cnt++;
}

bool spfa(int s)
{
    memset(vis, false, sizeof(vis));
    dis[s] = 0;
    queue <int> q;
    q.push(s);
    while(!q.empty())
    {
        int cur = q.front();
        q.pop();
        vis[cur] = false;
        for(int i = head2[cur]; i != -1; i = e2[i].next)
        {
            if(dis[e2[i].v] > dis[cur] + e2[i].w)
            {
                dis[e2[i].v] = dis[cur] + e2[i].w;
                if(!vis[e2[i].v])
                {
                    vis[e2[i].v] = true;
                    q.push(e2[i].v);
                }
            }
        }
    }
}

int A_star(int s, int t, int k)
{
    if(s == t) 
        k++;
    if(dis[s] == INF) 
        return -1;
    priority_queue <node1> q;
    int num = 0;
    node1 tmp, to;
    tmp.v = s;
    tmp.g = 0;
    tmp.f = tmp.g + dis[tmp.v];
    q.push(tmp);
    while(!q.empty())
    {
        tmp = q.top();
        q.pop();
        if(tmp.v == t) 
            num++;
        if(num == k) 
            return tmp.g;
        for(int i = head[tmp.v]; i != -1; i = e[i].next)
        {
            to.v = e[i].v;
            to.g = tmp.g + e[i].w;
            to.f = to.g + dis[to.v];
            q.push(to);
        }
    }
    return -1;
}

int main()
{
    scanf("%d %d", &n, &m);
    cnt = 0;
    memset(head, -1, sizeof(head));
    memset(head2, -1, sizeof(head2));
    for(int i = 0; i < MAX1; i++)
        dis[i] = INF;
    while(m--)
    {
        int u, v, w;
        scanf("%d %d %d", &u, &v, &w);
        Add(u, v, w);
    }
    int s, t, k;
    scanf("%d %d %d", &s, &t, &k);
    spfa(t);
    printf("%d\n", A_star(s, t, k));
}