LA 2666 最短路神题

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cmath>
#include<vector>
#include<set>
#include<algorithm>
#include<queue>
#define sqr(x) (x)*(x)
#define f1 first
#define f2 second
#define pb push_back
#define fr(i,x,y) for(int i=x;i<=y;++i)
#define ms(x,y) memset(x,y,sizeof(x))
#define sqr(x) (x)*(x)
typedef long long ll;
const int inf=0x3f3f3f3f;
#define N 30005
#define M 500005
using namespace std;
struct edge{
    int to,next,w;
}e[M];
struct node{
    int id;
    int d;
    node(int x=0,int y=0):id(x),d(y){}
    bool operator<(const node&oth)const{
        return d>oth.d;
    }
};
priority_queue<node> q;
int head[N],r[N];
int f[11][N];
int d[N];
bool vis[N];
int n,m,o;
void add(int x,int y,int z)
{
    e[o].to=y;
    e[o].w=z;
    e[o].next=head[x];
    head[x]=o++;
}
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
int dijkstra(int s)
{
    memset(d,0,(n+1)<<2);
    memset(vis,0,n+1);
    d[s]=0;
    while (!q.empty()) q.pop();
    q.push(node(s,d[s]));
    node tmp;int j,v;
    int re=0;
    while (!q.empty())
    {
        tmp=q.top();q.pop();
        if (vis[tmp.id])continue;
        j=tmp.id;
        vis[j]=1;re++;
        for (int k=head[j];k!=-1;k=e[k].next)
            {  v=e[k].to; 
                if (vis[v])continue;
                if (d[v]>d[j]+e[k].w||!d[v])
                    {   d[v]=d[j]+e[k].w;
                        if (f[r[s]+1][v]>d[v]) q.push(node(v,d[v]));
                    }
            }
    }
    //int re=0;fr(i,1,n)re+=(f[r[s]+1][i]>d[i]&&d[i]);return re+1;//O(n^2)致命TLE
    return re;
}

void bfs(int d[])
{   while (!q.empty()) q.pop();
    fr(i,1,n) if (!vis[i]){d[i]=inf;} else {d[i]=0;q.push(node(i,d[i]));vis[i]=0;}
    node tmp;
    int j,v;
    while (!q.empty())
    {
        tmp=q.top();q.pop();
        if (vis[tmp.id])continue;
        j=tmp.id;
        vis[j]=1;
        for (int k=head[j];k!=-1;k=e[k].next)
            {
                v=e[k].to;
                if (vis[v])continue;
                if (d[v]>d[j]+e[k].w)
                    {   d[v]=d[j]+e[k].w;
                        q.push(node(v,d[v]));
                    }
            }

    }
}
void doit()
{   scanf("%d%d",&n,&m);
    ms(head,255);o=0;
    int x,y,z;
    int ma=0;
    fr(i,1,n)
    {
        r[i]=read();//scanf("%d",&r[i]);
        ma=max(ma,r[i]);
    }
    fr(i,1,m)
    {   x=read();y=read();z=read();//scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }

    fr(i,1,ma)
    {   ms(vis,0);
        fr(j,1,n) if (r[j]>=i) vis[j]=1;
        bfs(f[i]);
    }

    int t1=0;
    fr(i,1,n) if (r[i]==ma) t1++;
    int ans=t1*n;

    fr(i,1,n)if (r[i]!=ma)
        {
            int tmp=dijkstra(i);
            ans+=tmp;
        }
    printf("%d\n",ans);
}
int main()
{   int cas;
    scanf("%d",&cas);
    int id=0;
    while (cas--)
         {
             if (id) puts("");id++;
             doit();
         }
}