洛谷-4716 【模板】最小树形图

题目描述
给定包含 n个结点， m 条有向边的一个图。试求一棵以结点 r为根的最小树形图，并输出最小树形图每条边的权值之和，如果没有以 r 为根的最小树形图，输出 -1。
输入格式
第一行包含三个整数 n,m,r，意义同题目所述。
接下来 m 行，每行包含三个整数 u,v,w，表示图中存在一条从 u 指向 v 的权值为 w 的有向边。
输出格式
如果原图中存在以 r 为根的最小树形图，就输出最小树形图每条边的权值之和，否则输出 -1。

输入输出样例
输入 #1
4 6 1
1 2 3
1 3 1
4 1 2
4 2 2
3 2 1
3 4 1

输出 #1
3

输入 #2
4 6 3
1 2 3
1 3 1
4 1 2
4 2 2
3 2 1
3 4 1

输出 #2
4

输入 #3
4 6 2
1 2 3
1 3 1
4 1 2
4 2 2
3 2 1
3 4 1

输出 #3
-1

说明/提示
数据范围
对于所有数据，$1\le u,v\le n\le 100,1\le m\le 10^4$

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 100 + 5;
const int M = 10000 + 5;
const ll INF = 0x3f3f3f3f;
int n, m, r;
struct edge {int u, v, w;} e[M];
int fa[N], id[N], top[N], minw[N];
ll get_ans (int n, int m) {
    ll ans = 0;
    while (true) {
        int cnt = 0;
        for (int i = 1; i <= n; ++i) {
            id[i] = top[i] = 0; minw[i] = INF;
        }
        for (int i = 0; i < m; ++i) {
            if (e[i].u != e[i].v && e[i].w < minw[e[i].v]) {
                fa[e[i].v] = e[i].u;
                minw[e[i].v] = e[i].w;
            }
        }
        minw[r] = 0;
        for (int i = 1; i <= n; ++i) {
            if (minw[i] == INF) return -1;
            ans += minw[i];
            int u = i;
            while (u != r && top[u] != i && !id[u]) {
                top[u] = i;
                u = fa[u];
            }
            if (u != r && !id[u]) {
                id[u] = ++cnt;
                for (int v = fa[u]; v != u; v = fa[v]) id[v] = cnt;
            }
        }
        if (cnt == 0) return ans;
        for (int i = 1; i <= n; ++i) {
            if (!id[i]) id[i] = ++cnt;
        }
        for (int i = 0; i < m; ++i) {
            int prew = minw[e[i].v];
            e[i].u = id[e[i].u];
            e[i].v = id[e[i].v];
            if (e[i].u != e[i].v) {
                e[i].w -= prew;
            }
        }
        n = cnt; r = id[r];
    }
}
int main () {
    cin >> n >> m >> r;
    for (int i = 0; i < m; ++i) {
        static int u, v, w;
        cin >> u >> v >> w; 
        e[i] = (edge) {u, v, w};
    }
    cout << get_ans (n, m) << endl;
}