CodeForces - 472D Design Tutorial: Inverse the Problem——最小生成树

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
typedef long long ll;
const int maxn = 2005;
struct Edge {
    int from, to, cost;
    bool operator < (const Edge &temp) const {
        return cost < temp.cost;
    }
}edge[maxn*maxn];
int n, a[maxn][maxn], tot, head[maxn];
struct E { int to, cost, next; } e[maxn*maxn];
void init() { tot = 0; memset(head, -1, sizeof(head)); }
void addedge(int u, int v, int cost) {
    e[++tot].to = v; e[tot].cost = cost; e[tot].next = head[u];
    head[u] = tot;
}

int par[maxn], ran[maxn];
int query(int x) { return par[x] == x ? x : par[x] = query(par[x]); }
bool unite(int x, int y) {
    x = query(x), y = query(y);
    if (x == y) return false;
    if (ran[x] < ran[y]) par[x] = y;
    else {
        par[y] = x;
        if (ran[x] == ran[y]) ran[x]++;
    }
    return true;
}

ll dis[maxn];
void dfs(int u, int f, ll d) {
    dis[u] = d;
    for (int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to, cost = e[i].cost;
        if (v == f) continue;
        dfs(v, u, d + cost);
    }
}

int iabs(int x) { return x > 0 ? x : -x; }

int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= n; j++) {
            scanf("%d", &a[i][j]);
        }
    }
    bool ok = true;
    int cnt = 0;
    for (int i = 1; i <= n; i++) {
        for (int j = i; j <= n; j++) {
            if (i == j && a[i][j] != 0) { ok = false; i = n + 1; break; }
            if (a[i][j] != a[j][i]) { ok = false; i = n + 1; break; }
            if (a[i][j] == 0) continue;
            edge[++cnt].from = i;
            edge[cnt].to = j;
            edge[cnt].cost = a[i][j];
        }
    }
    if (!ok) { printf("NO\n"); return 0; }
    sort(edge+1, edge+1+cnt);
    init();
    for (int i = 1; i <= n; i++) par[i] = i, ran[i] = 0;
    int num = 0;
    for (int i = 1; i <= cnt && num < n - 1; i++) {
        int u = edge[i].from, v = edge[i].to, cost = edge[i].cost;
        if (unite(u, v)) {
            ++num;
            addedge(u, v, cost); addedge(v, u, cost);
        }
    }
    if (num != n - 1) { printf("NO\n"); return 0; }
    for (int i = 1; i <= n; i++) {
        memset(dis, 0, sizeof(dis));
        dfs(i, -1, 0);
        for (int j = 1; j <= n; j++) {
            if (i == j) continue;
            if (dis[j] != a[i][j]) { ok = false; i = n + 1; break; }
        }
    }
    if (ok) printf("YES\n");
    else printf("NO\n");
    return 0;
}