FZU 2169 shadow spfa

题目链接：shadow

好佩服自己耶~~~好厉害~~~

麻麻再也不用担心我的spfa 和 邻接表技能了~~~

spfa 记录最短路径。

#include <stdio.h>
#include <string.h>
#include <iostream>
#include <queue>
#define maxn 100010
#define inf 1000000000
using namespace std;

int n, k;
int num[maxn];
int army[maxn];

struct Node {
    int v;
    int nxt;
}edge[maxn]; // 对于edge[i]这条边，（u-----v）只保存v和上一条以u为起点的边的序号。u的信息保存在最后一条以u为起点的边上。


bool visQue[maxn]; //顶点是否已经入队列
int cnt[maxn]; // 顶点入队列的次数
int dis[maxn]; // 保存源点到每个点的最短距离

int head[maxn];
int tot;
bool visArmy[maxn]; // 判断叛军是否已经被消灭
int fa[maxn]; // 最短路过程中记录每个节点的前驱

void addEdge(int u, int v) {
    edge[tot].v = v;
    edge[tot].nxt = head[u];
    head[u] = tot++;
}


bool spfa(int sou) {
    { // 初始化
        memset(visQue, 0, sizeof(visQue));
        memset(cnt, 0, sizeof(cnt));
        for (int i=0; i<maxn; ++i) {
            dis[i] = inf;
        }
    }

    dis[sou] = 0;
    queue<int>que;
    que.push(sou);
    visQue[sou] = 1;
    cnt[sou]++;

    while(!que.empty()) {
        int u  = que.front();
        que.pop();
        visQue[u] = 0;
        for (int i=u; i!=-1; i=edge[i].nxt) {
            int v = edge[i].v;
            if (dis[v] > dis[u] + 1) {
                dis[v] = dis[u] + 1;
                fa[v] = u;
                if (!visQue[v]) {  // 为什么在当前点被优化的时候才考虑是不是要把它加进队列呢？
                    que.push(v);  // 只有当前点被优化了，与它相邻的点的最短距离才有可能被修改，此时加进队列以松弛以它为顶点的边。
                    visQue[v] = 1;
                    cnt[v]++;
                    if (cnt[v] >= n) return false;
                }
            }
        }
    }
    return true;
}

int sum(int src) {
    int temp = 0;
    while (fa[src] != -1) {
        if (visArmy[src]) return temp;
        if (!visArmy[src]) {
            temp += num[src];
            visArmy[src] = 1;
        }
        src = fa[src];
    }
    return temp + num[1];
}


int main() {
    while(cin >> n >> k) {
        memset(head, -1, sizeof(head));
        tot = 0;
        memset(visArmy, 0, sizeof(visArmy));

        for (int i=0; i<n; ++i) { //输入每个城市的叛军数量
            cin >> num[i];
        }
        for (int i=0; i<k; ++i) {
            cin >> army[i]; // 输入有军队的城市编号。讲道理，开始还想着，遍历军队的时候也要从n个城市里面去找，蠢了。
        }
        for (int i=0; i<n-1; ++i) {
            int u, v;
            cin >> u >> v;
            addEdge(u, v);
        }

        fa[1] = -1;
        spfa(1); // 以1为源点找最短路
        int ans = 0;
        for (int i=0; i<k; ++i) {
            ans += sum(army[i]); // sum()求当前有军队的城市能消灭的叛军。
            //cout << sum(i) << "-----\n";
        }
        cout << ans << endl;
    }
    return 0;
}