CF1120D Power Tree

Solution Part1

子树操作，可以用DFS序解决。
这一题，我们实际上只用将叶子节点按DFS序从小到大排列，将其编号为1...n(n为叶子数量)。
对于点u，它控制的区间就是[cnt + 1,cnt + siz[u] + 1]，cnt表示比u的DFS序小的叶子的数量，siz[u]表示u的子节点中的叶子数量。
我们将这个序列差分，对于每个区间[a,b]，在点a与点b + 1之间连一条边，权值为控制[a,b]的代价，数字可以通过这条边流向另一个点，一端数字加w，另一端数字减w。
因为有的区间右端点可能为n，所以我们要多增加一个点n + 1，点n + 1的权值恒为0。这也意味着差分数组和为0。
这样我们的问题就变成了，在一个有n + 1个点的图上，选的边可以通过“搬运”两端的数字将图清零，求选边的最小代价。
很明显，只要图连通，数字就可以在图上自由流动，不管如何规定权值都可以将图清零。
求一遍最小生成树即可。

Solution Part2

第二行要我们输出可能存在于最优集合中的点，也就是求出所有最优集合的并集。
我们将其转化为求图上所有最小生成树的边集的并集。
我们考虑Kruskal算法。
Kruskal算法通过“合并”点集来求最小生成树。
我们将“合并”在一起的点通过类似强连通缩点的形式将其缩为一个点。此时图中所有点都属于不同的集合(因为缩点了)。若此时图中最小权值为w，此时所有权值为w的边都属于该图的某一个最小生成树。

Code

#include <bits/stdc++.h>

using namespace std;

const int MAXN = 2e5 + 10;

struct EDGE{
    int u,v,w,id;
}e[MAXN];

bool cmp(EDGE a,EDGE b) {
    return a.w < b.w;
}

vector<int> edge[MAXN];
int wei[MAXN],n,siz[MAXN],fa[MAXN],num[MAXN],cnt,lcnt,ncnt;
bool mark[MAXN];
long long ans;

int find(int x) {return x == fa[x]?x:fa[x] = find(fa[x]);}

void merge(int a,int b) {fa[find(a)] = find(b);return;}

void DFS(int u,int pre) {
    int len = edge[u].size();
    num[u] = ++cnt;
    if (len == 1 && u != 1) {
        lcnt++;
        e[num[u]] = {lcnt,lcnt + 1,wei[u],u};
        siz[u] = 1;
        return;
    }
    e[num[u]].u = lcnt + 1;
    for (int i = 0; i < len; i++) {
        int v = edge[u][i];
        if (v == pre) continue;
        DFS(v,u);
        siz[u] += siz[v];
    }
    e[num[u]].v = e[num[u]].u + siz[u];
    e[num[u]].w = wei[u];
    e[num[u]].id = u;
}

void Kruskal() {
    sort(e + 1,e + n + 1,cmp);
    for (int i = 1; i <= lcnt + 1; i++) fa[i] = i;
    int j = 1;
    for (int i = 1; i <= n; i = j + 1) {
        while (j + 1 <= n && e[j + 1].w == e[i].w) j++;
        for (int k = i; k <= j; k++) {
            if (find(e[k].u) != find(e[k].v)) mark[e[k].id] = true,ncnt++;
        }
        while (i <= j) {
            int a = e[i].u,b = e[i].v;
            if (find(a) != find(b)) {
                merge(a,b);
                ans += e[i].w;
            }
            i++;
        }
    }
    return;
}

int main() {
    scanf("%d",&n);
    for (int i = 1; i <= n; i++) {
        scanf("%d",&wei[i]);
    }
    for (int i = 1; i < n; i++) {
        int u,v;
        scanf("%d%d",&u,&v);
        edge[u].push_back(v);
        edge[v].push_back(u);
    }
    DFS(1,0);
    Kruskal();
    cout << ans << ' ' << ncnt << endl;
    for (int i = 1; i <= n; i++) {
        if (mark[i]) cout << i << ' ';
    }
    cout << endl;
    return 0;
}

注：判断叶子节点时也要判它是不是根，不然遇到只有两个点的树就挂了。