[数据结构与算法]317号子任务 CSP

#include <iostream>
#include <cstring>
#include <queue>
#include <algorithm>

using namespace std;

const int N = 10010,M = 20010,INF = 0x3f3f3f3f; // 无向边，边要取双倍

int n,m,k;
int dist[N]; // 排序后取前k个。  // 存储每个点到x号点的最短距离
int h[N], w[M], e[M], ne[M], idx; // next 开的数组也应该是M个，而不是N个 。超时就是因为这个，到官网一分都没拿到
int is_x[N];
bool st[N];
int ds[N][1010];
int len;

int u1,v1,w1;

void add(int a,int b,int c)
{
    e[idx] = b,w[idx] = c,ne[idx] = h[a],h[a] = idx++; // 边加错了！！！ c写成a了
}

void spfa(int k1)
{
    memset(dist, 0x3f, sizeof dist);
    dist[k1] = 0; // 第二位的k 会超过1000，所以要重新定义！
    st[N] = {0};
    queue<int> q;
    q.push(k1);
    st[k1] = true;

    while (q.size())
    {
        auto t = q.front();
        q.pop();

        st[t] = false;

        for (int i = h[t]; i != -1; i = ne[i])
        {
            int j = e[i];
            if (dist[j] > dist[t] + w[i])
            {
                dist[j] = dist[t] + w[i];
                if (!st[j])     // 如果队列中已存在j，则不需要将j重复插入
                {
                    q.push(j);
                    st[j] = true;
                }
            }
        }
    }
    
    for(int i=1;i <= n ;i++) ds[i][len] = dist[i];
}


int main()
{
    cin >> n >> m >> k;
    
    for(int i=1;i <= n;i++) 
    {
        cin >> is_x[i];
    }
    
    memset(h, -1, sizeof h); // 没有加这个导致超时。一定要记得初始化
    int tmp = m;
    while(tmp--)
    {
        cin >> u1 >> v1 >> w1;
        add(u1,v1,w1);
        add(v1,u1,w1);
    }
    
    
    for(int i=1;i <= n;i++)
    {
        if(is_x[i])
        {
            spfa(i);
            len++;
        }
    }
    
    for(int i=1;i <= n;i++)
    {
        int ans = 0;
        sort(ds[i],ds[i]+len);
        for(int j=0;j < len && j < k;j++)
        {
            if(ds[i][j] != INF)  ans += ds[i][j];
            else break;
        }
        cout << ans << endl;
    }
    
    
    return 0;
}