【题解】AT3611 Tree MST

喝，长大了🎵↑↓↑↓

思路

点分治 + kruskal.

对于完全图 MST 问题有一个结论：

$\forall G = (V, E)$，若 $E_1 \cup ... \cup E_m = E$，设 $E_i$ 经过 kruskal 得到的有效边集为 $E^{\prime}_i$，则对 $E^{\prime}_1 \cup ... \cup E^{\prime}_m$ 进行 kruskal 得到的是 $G$ 的 MST。

又发现题目考虑的是树上路径问题，所以想到点分治。

因为点分治可以考虑到树上的每一条路径，它们取并显然可以覆盖所有边，所以可以直接在点分的时候求当前子树的 MST。

因为点分的时候钦定子树重心为根，所以树上路径可以换一种形式写。

令 $dep_u$ 表示从当前分治的根结点到点 $u$ 的距离，那么在 $(u, v)$ 之间连边的代价是 $w_u + w_v + dep_u + dep_v$

可以把贡献拆开，$w_u + dep_u$ 赋给 $u$，$v$ 同理。

那么只需要在子树中找出 $w_u + dep_u$ 最小的点 $u$，然后子树内的其他所有点都向点 $u$ 连边即可。

那么只需要进行一遍点分就可以找出所有有用的点，直接 kruskal 一次就行。

点分的复杂度是 $O(n \log n)$，所以至多有 $O(n \log n)$ 条边，总复杂度 $O(n \log^2 n)$

代码

view code#include <cstdio>
#include <vector>
#include <algorithm>
using namespace std;

typedef long long ll;

const int maxn = 2e5 + 5;
const int maxm = 5e6 + 5;
const int inf = 0x3f3f3f3f;

struct item
{
    int u, v;
    ll w;

    bool operator < (const item& rhs) const { return (w < rhs.w); }
} edge[maxm];

struct node
{
    int p;
    ll w;

    bool operator < (const node& rhs) const { return (w < rhs.w); }
} nd[maxn];

int n, ecnt, len;
int rt, max_sz, tot_sz;
int w[maxn], fa[maxn], sz[maxn];
bool vis[maxn];
ll dep[maxn];
vector<int> g[maxn], l[maxn];

inline int read()
{
    int res = 0;
    char ch = getchar();
    while ((ch < '0') || (ch > '9')) ch = getchar();
    while ((ch >= '0') && (ch <= '9'))
    {
        res = res * 10 + ch - '0';
        ch = getchar();
    }
    return res;
}

inline void get_sz(int u, int fa)
{
    sz[u] = 1;
    for (int v : g[u])
    {
        if ((v == fa) || vis[v]) continue;
        get_sz(v, u);
        sz[u] += sz[v];
    }
}

inline void get_rt(int u, int fa)
{
    int res = 0;
    for (int v : g[u])
    {
        if ((v == fa) || vis[v]) continue;
        get_rt(v, u);
        res = max(res, sz[v]);
    }
    res = max(res, tot_sz - sz[u]);
    if (res < max_sz) rt = u, max_sz = res;
}

inline void get_dep(int u, int fa)
{
    nd[++len] = (node){u, w[u] + dep[u]};
    for (int i = 0; i < g[u].size(); i++)
    {
        int v = g[u][i], d = l[u][i];
        if ((v == fa) || vis[v]) continue;
        dep[v] = dep[u] + d;
        get_dep(v, u);
    }
}

inline void solve(int u)
{
    // printf("solving %d\n", u);
    vis[u] = true;
    nd[len = 1] = (node){u, w[u]};
    for (int i = 0; i < g[u].size(); i++)
    {
        int v = g[u][i], d = l[u][i];
        if (vis[v]) continue;
        dep[v] = d;
        get_dep(v, u);
    }
    sort(nd + 1, nd + len + 1);
    // printf("debug %d\n", dep[nd[2].p]);
    for (int i = 2; i <= len; i++) edge[++ecnt] = (item){nd[i].p, nd[1].p, nd[i].w + nd[1].w};
    for (int v : g[u])
    {
        if (vis[v]) continue;
        get_sz(v, 0);
        tot_sz = sz[v], max_sz = inf;
        get_rt(v, 0);
        solve(rt);
    }
}

inline int get(int u) { return (fa[u] == u ? u : fa[u] = get(fa[u])); }

inline void kruskal()
{
    ll ans = 0;
    sort(edge + 1, edge + ecnt + 1);
    for (int i = 1; i <= n; i++) fa[i] = i;
    for (int i = 1; i <= ecnt; i++)
    {
        int fu = get(edge[i].u), fv = get(edge[i].v);
        if (fu != fv) { fa[fu] = fv, ans += edge[i].w; }
    }
    printf("%lld\n", ans);
}

int main()
{
    n = read();
    for (int i = 1; i <= n; i++) w[i] = read();
    for (int i = 1, u, v, d; i <= n - 1; i++)
    {
        u = read(), v = read(), d = read();
        g[u].push_back(v), l[u].push_back(d);
        g[v].push_back(u), l[v].push_back(d);
    }
    tot_sz = n, max_sz = inf;
    get_sz(1, 0), get_rt(1, 0);
    solve(rt);
    kruskal();
    return 0;
}

/*
4
1 3 5 1
1 2 1
2 3 2
3 4 3

4 <-> 1 : 8
3 <-> 1 : 9
2 <-> 1 : 8
2 <-> 1 : 11

1 <-> 1 : 4
2 <-> 1 : 5
3 <-> 1 : 9
4 <-> 1 : 8
1 <-> 1 : 2
3 <-> 4 : 9
4 <-> 4 : 8
4 <-> 4 : 2
*/