Power Tree

GitHub跳转原题关系图返回列表

把子树加值改写成差分数组上的两次点修改，购买集合可行等价于对应边让 0..m 连通，于是第一问变成最小生成树；第二问用按权值分组的 Kruskal 判定每条边是否属于某棵 MST。

OJ: codeforces

题目 ID: 1120D

难度:提高+/省选-

标签:最小生成树并查集树上问题DFS序python

创建: 2026-10-21 09:10

更新: 2026-10-28 17:33

目录

  1. 形式化题目
  2. 暴力解法
    1. 思路
    2. 代码
    3. 复杂度
    4. 瓶颈
  3. 正解
    1. 思路
    2. 代码
    3. 复杂度
  4. 总结

形式化题目

给定一棵以 11 为根、共 nn 个点的有根树，每个点有一个非负代价 cvc_v。叶子指「非根且度数为 11」的点，设叶子共 mm 个。

先选一个非空的点集 SS（代价为 ∑v∈Scv\sum_{v\in S} c_v），之后对手会给每个叶子填入任意整数。要求：无论对手填什么，都必须存在一组整数 xv (v∈S)x_v\ (v\in S)，使得对每个叶子 aa，落在 aa 上的总增量为

∑v∈Sa 在 v 的子树内xv  =  − (对手在 a 上填的数). \sum_{\substack{v\in S\\ a \text{ 在 } v \text{ 的子树内}}} x_v \;=\; -\,(\text{对手在 } a \text{ 上填的数}) .

第一问问满足要求的最小代价 ss；第二问求所有「存在某个代价为 ss 的可行集合包含它」的点。

暴力解法

思路

题目里只有「子树内叶子」这一种结构，所以最直接的做法就是照定义暴力：把 mm 个叶子按从根出发的 DFS 序编号成 1..m1..m，记 lov,hivlo_v,hi_v 为点 vv 子树内叶子的最小编号与最大编号。因为子树在 DFS 序里是连续的一段，vv 的子树叶子恰好是区间 [lov,hiv][lo_v,hi_v]——这是后面一切推导的地基。

然后枚举全部 2n2^n 个非空点集 SS。判定 SS 是否可行，可以不去碰线性代数，而是用一个组合判据（等价性在正解一节证明）：把每个点 vv 看成一条连接 lov−1lo_v-1 与 hivhi_v 的边，这 nn 条边落在 m+1m+1 个点（编号 0..m0..m）上；SS 可行当且仅当 SS 里的边能让这 m+1m+1 个点全部连通。

判连通只需一个并查集。把可行集合里代价最小的记为 ss，再把所有达到 ss 的可行集合取并集、升序输出，就是第二问答案。

复杂度 O(2n⋅(n+m))O(2^n \cdot (n+m))，n⩽8n \leqslant 8 时约 255×16255 \times 16 次操作，正好用来给正解当对拍答案。它同时也把题意验证清楚了：这个暴力完全不涉及生成树，只用题面定义加一个连通判据。

代码

cpp
/**
 * Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
 * rbook: -> https://rbook.roj.ac.cn  https://rbook2.roj.ac.cn
 * rainboy的学习导航网站: https://idx.roj.ac.cn
 * create_at: 2026-10-21 09:10
 * update_at: 2026-10-21 09:10
 */
// brute.cpp：n <= 8 的小数据暴力解，用来理解题意并给 main.cpp 当对拍答案。
//
// 暴力思路（完全不涉及 MST，只按题面定义验证）：
//   枚举所有非空的购买集合 S；
//   把叶子按 DFS 序编号 1..m，点 v 的子树叶子是一段区间 [lo_v, hi_v]，
//   对应新图 m+1 个点（0..m）上的边 (lo_v - 1, hi_v)；
//   检查 S 对应的这些边能否让 0..m 全部连通——连通即说明能清零任意叶子向量（题解里有论证）；
//   取连通方案里花费最小者作为第一问答案；
//   再把所有最优集合里的点求并集，升序输出，作为第二问答案。
//
// 复杂度 O(2^n * (n + m))，n <= 8 时约 255 * 16 次操作，随便跑。
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 15;   // n <= 8，留足余量

int n;
int c[MAXN];
vector<int> adj[MAXN];

int par[MAXN];
int lo[MAXN], hi[MAXN];
int m;                 // 叶子个数

// ---------- 并查集：只用来判断 0..m 是否连通 ----------
int fa[MAXN];
int find(int x) {
    while (fa[x] != x) { fa[x] = fa[fa[x]]; x = fa[x]; }
    return x;
}

// 为一次暴力判定准备叶子编号与区间（树是固定的，只需算一次，这里写成函数便于阅读）
void build() {
    // 迭代 DFS 求前序（n <= 8 递归也不会爆栈，但和 main.cpp 保持同一种写法）
    vector<int> ord;
    vector<char> vis(n + 1, 0);
    vector<int> st;
    st.push_back(1);
    vis[1] = 1;
    par[1] = 0;
    while (!st.empty()) {
        int u = st.back();
        st.pop_back();
        ord.push_back(u);
        for (int v : adj[u]) {
            if (!vis[v]) { vis[v] = 1; par[v] = u; st.push_back(v); }
        }
    }

    for (int v = 1; v <= n; ++v) { lo[v] = INT_MAX; hi[v] = -1; }
    m = 0;
    for (int u : ord) {
        // 叶子 = 非根、度数为 1 的点
        if (u != 1 && (int)adj[u].size() == 1) { ++m; lo[u] = hi[u] = m; }
    }
    for (int i = (int)ord.size() - 1; i >= 0; --i) {
        int u = ord[i];
        if (par[u] != 0) {
            int p = par[u];
            lo[p] = min(lo[p], lo[u]);
            hi[p] = max(hi[p], hi[u]);
        }
    }
}

// 检查购买集合（mask 的第 (v-1) 位为 1 表示买了 v）能否清零任意叶子向量
bool feasible(int mask) {
    for (int i = 0; i <= m; ++i) fa[i] = i;
    for (int v = 1; v <= n; ++v) {
        if (!((mask >> (v - 1)) & 1)) continue;
        int a = find(lo[v] - 1), b = find(hi[v]);
        if (a != b) fa[a] = b;
    }
    int r = find(0);
    for (int i = 1; i <= m; ++i) {
        if (find(i) != r) return false;
    }
    return true;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;
    for (int i = 1; i <= n; ++i) cin >> c[i];
    for (int i = 1; i < n; ++i) {
        int a, b;
        cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    build();

    long long best = LLONG_MAX;
    // 第一轮：找出最小花费
    for (int mask = 1; mask < (1 << n); ++mask) {
        if (!feasible(mask)) continue;
        long long cost = 0;
        for (int v = 1; v <= n; ++v) {
            if ((mask >> (v - 1)) & 1) cost += c[v];
        }
        best = min(best, cost);
    }

    // 第二轮：把所有达到最小花费的集合里的点求并集
    vector<char> inAnswer(n + 1, 0);
    for (int mask = 1; mask < (1 << n); ++mask) {
        if (!feasible(mask)) continue;
        long long cost = 0;
        for (int v = 1; v <= n; ++v) {
            if ((mask >> (v - 1)) & 1) cost += c[v];
        }
        if (cost != best) continue;
        for (int v = 1; v <= n; ++v) {
            if ((mask >> (v - 1)) & 1) inAnswer[v] = 1;
        }
    }

    vector<int> answer;
    for (int v = 1; v <= n; ++v) {
        if (inAnswer[v]) answer.push_back(v);
    }

    cout << best << ' ' << answer.size() << '\n';
    for (size_t i = 0; i < answer.size(); ++i) {
        cout << answer[i] << (i + 1 == answer.size() ? '\n' : ' ');
    }
    return 0;
}

复杂度

时间 O(2n⋅(n+m))O(2^n\cdot(n+m))，空间 O(n)O(n)。

瓶颈

暴力暴露了两处浪费，而且它们是同一个问题的两面：

  • 枚举层面，2n2^n 个集合里绝大多数是「明显拼凑」的，我们却在逐个重新判连通。说明可行集合的结构还没被利用——它其实是一个「连通性 + 最小权」的组合对象，天生该用最小生成树一次求出。
  • 判定层面，暴力的连通判据是「凭空给出」的，没有解释为什么「边连通 ⇔\Leftrightarrow 能清零任意叶子填充」。不知道等价性来自哪里，就无从下手优化。

正解逐个解决：先把「子树加值」翻译成差分数组上的点修改，得到每点一条边的建模；等价性由差分向量的秩给出；最后最小权连通边集就是最小生成树。

正解

思路

第一步：子树加值 = 差分数组上的两次点修改。

先把叶子按 DFS 序编号 1..m1..m。对任意点 vv，它的子树叶子在 DFS 序中连续，记作区间 [lov,hiv][lo_v,hi_v]。

现在把「叶子上的取值数组」换成它的差分「首尾闭合」形式：对叶子向量 y1..ymy_1..y_m（只有叶子有值，叶子之外视作 00），定义长度为 m+1m+1 的数组 d0..dmd_0..d_m：

d0=y1,dj=yj+1−yj (1⩽j⩽m−1),dm=− ym. d_0 = y_1,\qquad d_j = y_{j+1}-y_j\ (1\leqslant j\leqslant m-1),\qquad d_m = -\,y_m .

y↦dy \mapsto d 是 mm 维空间到「和为 00 的 m+1m+1 维向量」的双射：∑j=0mdj=0\sum_{j=0}^{m} d_j = 0 且 yy 可由前缀和还原。反过来，一次「把区间 [l,r][l,r] 内所有叶子同时加 xx」在这组坐标里只改动两个位置：

dl−1 += x,dr −= x. d_{l-1}\ {+}{=}\ x,\qquad d_{r}\ {-}{=}\ x .

这就是全部建模的起点：子树区间加值 ⟶\longrightarrow 差分数组上「一端 +x+x、另一端 −x-x」的一次点对修改。

第二步：每个顶点对应新图上的一条边。

把 dd 的 m+1m+1 个下标 0,1,…,m0,1,\dots,m 看成新图的 m+1m+1 个顶点。点 vv 的子树区间是 [lov,hiv][lo_v,hi_v]，它对差分的作用就是

dlov−1 += x,dhiv −= x, d_{lo_v-1}\ {+}{=}\ x,\qquad d_{hi_v}\ {-}{=}\ x ,

也就是说，点 vv 对应新图上的一条边 (lov−1, hiv)(lo_v-1,\ hi_v)，边权取 cvc_v。至此每个顶点恰好贡献一条边，问题从树上搬到了这张 m+1m+1 个点、nn 条边的图上。

顺带看一眼图的连通性：叶子 v=kv=k 的区间是 [k,k][k,k]，给出边 (k−1,k)(k-1,k)。这 mm 条边连成 0−1−2−⋯−m0-1-2-\cdots-m 的链，所以新图必然连通，最小生成树一定存在，答案一定有定义。

第三步：可行 ⇔\Leftrightarrow 对应的边集让 0..m0..m 全连通。

对手可以填任意整数向量，所以要求是：这些操作能生成任意叶子向量，也就是任意满足 ∑jdj=0\sum_j d_j=0 的差分向量。

给点 vv 配上它的系数向量 eve_v（dlov−1d_{lo_v-1} 处是 +1+1、dhivd_{hi_v} 处是 −1-1，其余为 00），定义域取在 ∑jdj=0\sum_j d_j=0 这个超平面上。那么

购买集合 S 可行  ⟺  { ev:v∈S } 张成整个 m 维空间. \text{购买集合 } S \text{ 可行} \iff \{\,e_v : v\in S\,\} \text{ 张成整个 } m \text{ 维空间}.

现在关键的一步：把所有 eve_v 看成图 GSG_S（顶点 0..m0..m，边集为 SS 对应的那 ∣S∣|S| 条边）的有向关联矩阵的列。线性代数给出，图的关联矩阵的秩等于

(m+1)−c, (m+1) - c,

其中 cc 是 GSG_S 的连通块数。而每个 eve_v 各分量之和为 00，张成空间落在那个 mm 维超平面内，所以「张成整个 mm 维空间」当且仅当秩恰为 mm，即 c=1c=1：

{ ev:v∈S } 张成 m 维空间  ⟺  GS 连通（0..m 全在一个连通块里）. \{\,e_v : v\in S\,\} \text{ 张成 } m \text{ 维空间} \iff G_S \text{ 连通（}0..m \text{ 全在一个连通块里）}.

可逆性（整数解）也不成问题：取 GSG_S 的一棵生成树，按自底向上的顺序处理每个非根顶点 uu——它的唯一父边只影响 uu 与父点，此时 uu 的目标值已完全确定，据此定出这条边对应的 xx（整数减整数仍是整数）；最后根 00 因为 ∑jdj=0\sum_j d_j=0 自动吻合。所以「连通」给出的不只是实数解，而是整数解，这正是题目需要的。

第四步：第一问 = 最小生成树权值。

由第三步，SS 可行   ⟺  S\iff S 对应的边集连通 0..m0..m。任何连通的边集都包含一棵生成树，而边权 cv⩾0c_v \geqslant 0，删掉多余边只会让代价不增，所以

min⁡S 可行∑v∈Scv  =  min⁡生成树 T∑e∈Twe, \min_{S \text{ 可行}} \sum_{v\in S} c_v \;=\; \min_{\text{生成树 } T} \sum_{e\in T} w_e ,

即新图的最小生成树权值。用 Kruskal 排序后一遍并查集即可求出，total 累加真正被采纳的边的权值。

样例 1 的建模结果。叶子是 3,4,53,4,5，按本实现的 DFS 序它们依次被编成 5o15 o1、4o24 o2、3o33 o3，故 m=3m=3，新图顶点为 0..30..3：

顶点 vv 子树内叶子 叶子区间 新图边 边权
1 3,4,5 [1,3] (0,3) 5
2 3,4 [2,3] (1,3) 1
3 3 [3,3] (2,3) 3
4 4 [2,2] (1,2) 2
5 5 [1,1] (0,1) 1

最小生成树取边 (0,1)(0,1)、(1,2)(1,2)、(1,3)(1,3)，对应顶点 5,4,25,4,2，合计 1+2+1=41+2+1=4，与题面第一问一致。而 2,4,52,4,5 三个点落在某棵 MST 里，正是题面第二问的答案；其余重边如 (0,3)(0,3)、(2,3)(2,3) 在只考虑更小权值边时两端已连通，不属于任何 MST。

第五步：第二问——边是否属于「某棵」MST。

需要一个不重建并查集就能批量判定的判据：

边 e 属于某棵 MST  ⟺  只用权值严格小于 ce 的边时，e 的两端仍不连通. \text{边 } e \text{ 属于某棵 MST} \iff \text{只用权值严格小于 } c_e \text{ 的边时，}e \text{ 的两端仍不连通}.

两个方向都要成立：

  • 判据成立 ⇒\Rightarrow 存在 MST 含 ee：先只用更小的边跑 Kruskal，此时两端分处两个连通块，把 ee 加进去是合法的（不成环）；继续按序跑完整套 Kruskal，就得到一棵包含 ee 的 MST。
  • ee 在环上是「严格最重边」⇒\Rightarrow 任何 MST 都不含 ee：若只用更小的边两端已连通，那么在这条全由更小边组成的路径加 ee 构成一个环，ee 是该环上唯一可能的最大权边（其余都严格更小）。由环性质，把 ee 换成环上另一条边不会让总权变大，所以总存在不含 ee 的 MST，也就没有 MST 必须含 ee。

于是判定可以按权值分组完成，这是本题实现的关键：把所有边按权值排序后逐组处理，处理到权值 ww 这一组时，并查集里恰好是「只用权值 <w< w 的边」的连通状态——这就是本组的判定基线。于是每组做两遍循环：

  1. 先对组内每条边查一次基线：两端不连通就说明它属于某棵 MST，收进答案；
  2. 再统一把整组边合并进并查集，得到下一组要用的基线。

必须先判完整组再合并：同权边之间会互相「接通」，如果边判边合，组内靠前被合并的边会污染后面兄弟边的判定，把本该入选的边误判掉。这也是本题最常见的错误来源。分组做法总代价 O(n α(n))O(n\,\alpha(n))，避免了对每条边重建并查集的 O(n2)O(n^2)。

对拍时专门用大量重复权值的数据（生成器里所有叶子代价只取 0,1,20,1,2 的模式）压这一点：权值全相等时第二问的答案退化成「整张图的桥」，最容易暴露先判后合写错的实现。

代码

C++ 版（Kruskal 按权值分组，一遍同时算出两问）：

cpp
/**
 * Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
 * rbook: -> https://rbook.roj.ac.cn  https://rbook2.roj.ac.cn
 * rainboy的学习导航网站: https://idx.roj.ac.cn
 * create_at: 2026-10-21 09:10
 * update_at: 2026-10-21 09:10
 */
// CF1120D Power Tree
// 题意：给一棵以 1 为根的树，每个点有权值 c_v。买一些点，买完之后对任意叶子取值，
//       都能通过「给某个已买点的子树内所有叶子同时加同一个整数」把叶子清零。
//       第一问：最小花费；第二问：哪些点出现在某个最优购买集合里。
//
// ===== 关键建模（本文件的核心，全部推导浓缩在下面几行）=====
// 把叶子按 DFS 序编号 1..m。点 v 的子树内叶子恰好是一段连续区间 [lo_v, hi_v]。
// 「给 v 的子树叶子整体 +x」 等价于在一维差分数组 d[0..m] 上做两次点修改：
//     d[lo_v - 1] += x ,  d[hi_v] -= x
// 于是每个点 v 唯一对应新图上的边 (lo_v - 1, hi_v)，边权 c_v；新图共 m+1 个点（编号 0..m）。
//
// 第一问 = 这张新图的最小生成树权值。为什么？
//   * 购买集合 S 能清零任意叶子向量 <=> {差分向量 of v | v in S} 张成整个 m 维空间
//     （叶子向量与差分向量一一对应，后者有 m 个自由度且和为 0 的约束由 d[0], d[m] 抵消）。
//   * 每个向量都是 (+1 在 lo-1) + (-1 在 hi) 的形式，恰好是一条「边上两端一加一减」的模式；
//     这样一族向量的秩为 m（满秩）当且仅当它们对应的边让 0..m 这 m+1 个点全部连通。
//     （连通 -> 沿任一生成树自底向上可依次独立调整每个点；不连通 -> 某个连通块内的差分和为 0，
//       存在张成不到的向量。）
//   * 「边集让 0..m 连通」当且仅当它包含一棵生成树，因此取权值最小的生成树最省。
//   顺带一提：所有叶子点 v 的区间是 [k,k]，给出边 (k-1, k)，所以新图必然连通（一条链），MST 一定存在。
//
// 第二问 = 判定每条边是否属于「某棵」MST，等价判定：把所有权值严格小于 c_e 的边加进来时，
//   e 的两个端点是否还不连通。实现上就是分组 Kruskal：
//   处理到权值 w 这一组时，并查集里恰好是「只用 < w 的边」的连通状态，
//   先拿这个基线状态对组内所有边判一次，再统一把整组合并掉。
//   这样做避免了「对每条边重建一次并查集」的 O(n^2)。
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

// n <= 2e5
const int MAXN = 200005;

int n;
int c[MAXN];                 // c[v]：买下点 v 的花费
vector<int> adj[MAXN];       // 原树邻接表

int par[MAXN];               // 迭代 DFS 得到的父亲
int lo[MAXN], hi[MAXN];      // 点 v 的子树内叶子编号区间 [lo, hi]

// 新图上的一条边：连接 u、v，权 w，来源顶点 id
struct Edge {
    int u, v, w, id;
    bool operator<(const Edge &o) const { return w < o.w; }
};
vector<Edge> edge;

// ---- 并查集（路径压缩 + 按大小合并），本题只需按大小合并即可保证复杂度 ----
int fa[MAXN], sz[MAXN];

void dsu_init(int n) {
    for (int i = 0; i <= n; ++i) { fa[i] = i; sz[i] = 1; }
}

int find(int x) {
    while (fa[x] != x) { fa[x] = fa[fa[x]]; x = fa[x]; }  // 路径减半，迭代写法不递归
    return x;
}

// 合并成功返回 true（这条边被 MST 采纳）
bool unite(int a, int b) {
    a = find(a); b = find(b);
    if (a == b) return false;
    if (sz[a] < sz[b]) swap(a, b);
    fa[b] = a; sz[a] += sz[b];
    return true;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;
    for (int i = 1; i <= n; ++i) cin >> c[i];
    for (int i = 1; i < n; ++i) {
        int a, b;
        cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    // ---------- 1. 显式栈求 DFS 序（n = 2e5 的链会爆递归栈，必须迭代） ----------
    // 栈式 DFS：弹出 u 时把未访问的儿子全部压栈，之后最晚压入的儿子最先被完整展开，
    // 因此得到的 preorder 是真正的 DFS 前序，每棵子树在其中是连续一段，叶子也按 DFS 序出现。
    vector<int> ord;
    ord.reserve(n);
    {
        vector<char> vis(n + 1, 0);
        vector<int> st;
        st.reserve(n);
        st.push_back(1);
        vis[1] = 1;
        par[1] = 0;
        while (!st.empty()) {
            int u = st.back();
            st.pop_back();
            ord.push_back(u);
            for (int v : adj[u]) {
                if (!vis[v]) {
                    vis[v] = 1;
                    par[v] = u;
                    st.push_back(v);
                }
            }
        }
    }

    // ---------- 2. 叶子按 DFS 序编号 1..m ----------
    // 题面定义：叶子 = 非根、度数为 1 的点。
    int m = 0;
    for (int v = 1; v <= n; ++v) { lo[v] = INT_MAX; hi[v] = -1; }
    for (int u : ord) {
        if (u != 1 && (int)adj[u].size() == 1) {   // u 是叶子
            ++m;
            lo[u] = hi[u] = m;
        }
    }

    // 逆着 DFS 序回溯，把子树的叶子区间并到父亲：O(n)
    for (int i = (int)ord.size() - 1; i >= 0; --i) {
        int u = ord[i];
        if (par[u] != 0) {
            int p = par[u];
            if (lo[u] < lo[p]) lo[p] = lo[u];
            if (hi[u] > hi[p]) hi[p] = hi[u];
        }
    }

    // ---------- 3. 建新图：m+1 个点（0..m），点 v 对应边 (lo_v - 1, hi_v) ----------
    edge.reserve(n);
    for (int v = 1; v <= n; ++v) {
        edge.push_back(Edge{lo[v] - 1, hi[v], c[v], v});
    }
    sort(edge.begin(), edge.end());

    // ---------- 4. 按权值分组的 Kruskal，同时回答两问 ----------
    dsu_init(m);
    ll total = 0;                 // 第一问：MST 总权值
    vector<int> answer;           // 第二问：属于某棵 MST 的边对应的点
    answer.reserve(n);

    for (int i = 0; i < n; ) {
        int j = i;
        while (j < n && edge[j].w == edge[i].w) ++j;   // [i, j) 是本组同权边

        // (a) 先用「只看 < w 的边」的并查集基线判断组内每条边
        //     两端在基线里不连通 <=> 存在某棵 MST 含这条边
        for (int k = i; k < j; ++k) {
            if (find(edge[k].u) != find(edge[k].v)) answer.push_back(edge[k].id);
        }
        // (b) 再把整组统一合并（同一组内先判后合，避免组内边互相影响判断）
        for (int k = i; k < j; ++k) {
            if (unite(edge[k].u, edge[k].v)) total += edge[k].w;
        }
        i = j;
    }

    sort(answer.begin(), answer.end());

    cout << total << ' ' << answer.size() << '\n';
    for (size_t i = 0; i < answer.size(); ++i) {
        cout << answer[i] << (i + 1 == answer.size() ? '\n' : ' ');
    }
    return 0;
}

Python 版（同一算法；DFS 用显式栈，邻接表用 CSR 数组把 n=2×105n=2\times10^5 的内存压到约 40MB）：

python
#!/usr/bin/env python3
# Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
# rbook: -> https://rbook.roj.ac.cn  https://rbook2.roj.ac.cn
# rainboy的学习导航网站: https://idx.roj.ac.cn
# create_at: 2026-10-21 09:10
# update_at: 2026-10-21 09:10
#
# main.py：CF1120D Power Tree，与 main.cpp 同一算法。
#
# 建模：叶子按 DFS 序编号 1..m，点 v 的子树叶子是一段连续区间 [lo_v, hi_v]。
#   「给 v 的子树叶子整体 +x」 等价于一维差分数组 d[0..m] 上的两次点修改：
#       d[lo_v - 1] += x ,  d[hi_v] -= x
#   于是每个点 v 对应新图 m+1 个点（编号 0..m）上的一条边 (lo_v-1, hi_v)，边权 c_v。
#   第一问 = 该图的最小生成树权值；第二问 = 属于某棵 MST 的边所对应的点。
# DFS 用显式栈（n 可达 2e5，Python 递归必爆），邻接表用 CSR 数组压内存。

import sys
from array import array


def main() -> None:
    tokens = sys.stdin.buffer.read().split()
    it = iter(tokens)
    n = int(next(it))

    # cost[v]：买下点 v 的花费。用 array('q') 存，避免 20 万个 int 对象
    cost = array("q", bytes(8 * (n + 1)))
    for v in range(1, n + 1):
        cost[v] = int(next(it))

    # ---------- CSR 邻接表：head/next 两条数组，比 list-of-list 省一半以上内存 ----------
    eu = array("i", bytes(4 * (n - 1)))
    ev = array("i", bytes(4 * (n - 1)))
    deg = array("i", bytes(4 * (n + 1)))
    for i in range(n - 1):
        a = int(next(it))
        b = int(next(it))
        eu[i] = a
        ev[i] = b
        deg[a] += 1
        deg[b] += 1
    del tokens, it  # 立刻释放 60 万个 bytes 对象，峰值内存能降一大截

    start = array("i", bytes(4 * (n + 2)))
    for v in range(1, n + 2):
        start[v] = start[v - 1] + deg[v - 1]
    pos = array("i", start)          # 每条边在 CSR 里的写入位置
    to = array("i", bytes(4 * 2 * (n - 1)))
    for i in range(n - 1):
        a, b = eu[i], ev[i]
        to[pos[a]] = b
        pos[a] += 1
        to[pos[b]] = a
        pos[b] += 1
    del eu, ev, deg, pos

    # ---------- 1. 显式栈求 DFS 前序，同时记录父亲 ----------
    # 弹栈时把儿子全部压栈：「最晚压入的最先展开」，所以得到的仍是合法 DFS 前序，
    # 每棵子树在其中连续，叶子也按 DFS 序出现。
    parent = array("i", bytes(4 * (n + 1)))
    visited = bytearray(n + 1)
    order = array("i")
    stack = [1]
    visited[1] = 1
    while stack:
        u = stack.pop()
        order.append(u)
        for k in range(start[u], start[u + 1]):
            v = to[k]
            if not visited[v]:
                visited[v] = 1
                parent[v] = u
                stack.append(v)
    del stack, visited, start, to

    # ---------- 2. 叶子按 DFS 序编号 1..m ----------
    # 叶子 = 非根、度数为 1 的点（题面定义）
    INF = n + 1
    lo = array("i", [INF]) * (n + 1)   # 子树内最小叶子编号
    hi = array("i", [-1]) * (n + 1)    # 子树内最大叶子编号
    # 用「谁有儿子」标记来判叶子：没有儿子且不是根 <=> 度数为 1 的非根点
    m = 0
    has_child = bytearray(n + 1)
    for u in order:
        p = parent[u]
        if p:
            has_child[p] = 1
    for u in order:
        if u != 1 and not has_child[u]:
            m += 1
            lo[u] = m
            hi[u] = m
    del has_child

    # 逆 DFS 序回溯，把儿子的区间并给父亲，O(n)
    for i in range(n - 1, -1, -1):
        u = order[i]
        p = parent[u]
        if p:
            if lo[u] < lo[p]:
                lo[p] = lo[u]
            if hi[u] > hi[p]:
                hi[p] = hi[u]

    # ---------- 3. 新图有 m+1 个点（0..m），点 v 对应边 (lo[v]-1, hi[v]) ----------
    # 按权值排序点编号即可代表按边权排序；Python 排序稳定，同权边自然聚在相邻位置。
    order_v = sorted(range(1, n + 1), key=cost.__getitem__)

    # ---------- 4. 按权值分组的 Kruskal，同时得出两问 ----------
    dsu = list(range(m + 1))
    size = [1] * (m + 1)

    def find(x: int) -> int:
        # 迭代路径压缩，避免递归深度问题
        root = x
        while dsu[root] != root:
            root = dsu[root]
        while dsu[x] != root:
            dsu[x], x = root, dsu[x]
        return root

    total = 0       # 第一问：MST 总权值
    answer = []     # 第二问：出现在某棵 MST 里的边对应的点

    i = 0
    while i < n:
        w = cost[order_v[i]]
        j = i
        while j < n and cost[order_v[j]] == w:
            j += 1

        # (a) 此时并查集恰好是「只用权值 < w 的边」的连通状态（基线）：
        #     两端在基线里不连通 <=> 这条边属于某棵 MST
        for k in range(i, j):
            v = order_v[k]
            if find(lo[v] - 1) != find(hi[v]):
                answer.append(v)
        # (b) 再统一合并整组，保证组内同权边不会影响彼此的判定
        for k in range(i, j):
            v = order_v[k]
            ru, rv = find(lo[v] - 1), find(hi[v])
            if ru != rv:
                if size[ru] < size[rv]:
                    ru, rv = rv, ru
                dsu[rv] = ru
                size[ru] += size[rv]
                total += w
        i = j

    answer.sort()
    out = ["%d %d" % (total, len(answer)), " ".join(map(str, answer))]
    sys.stdout.write("\n".join(out) + "\n")


if __name__ == "__main__":
    main()

复杂度

  • 求 DFS 序与 lo/hilo/hi：显式栈一遍 O(n)O(n)；因为 nn 可达 2×1052\times 10^5，深链会爆递归栈，两版实现都刻意写成迭代。
  • 排序（Kruskal 前按边权排序）：O(nlog⁡n)O(n\log n)。
  • 第一问 Kruskal 主体：O(n α(n))O(n\,\alpha(n))。
  • 第二问按权值分组判定：每组两遍线性扫描，总 O(n α(n))O(n\,\alpha(n))，不是 O(n2)O(n^2)。
  • 总体时间 O(nlog⁡n)O(n\log n)（排序是瓶颈），空间 O(n)O(n)。

总结

  • 全部内容起源于一次改写：子树加值 = 一维差分上「lov−1lo_v-1 处 +x+x、hivhi_v 处 −x-x」的两次点修改。叶子按 DFS 序编号后子树叶子是连续区间，这条改写才成立。
  • 于是每个顶点变成新图 m+1m+1 个点（0..m0..m）上的一条边 (lov−1,hiv)(lo_v-1,hi_v)，边权 cvc_v；「能清零任意叶子填充」等价于「这些差分向量张成整个 mm 维空间」，而由关联矩阵秩 =(m+1)−=(m+1)- 连通块数，等价于「对应边集让 0..m0..m 全连通」。
  • 连通的边集必含生成树、边权又非负，所以第一问就是新图的最小生成树权值。
  • 第二问用 MST 的经典边判定：只用严格更小的边时两端不连通   ⟺  \iff 属于某棵 MST。实现上按权值分组做 Kruskal，先判完整组再统一合并，避免 O(n2)O(n^2) 的重建并查集；忘记「先判后合」会让同权边互相污染，是本题最隐蔽的错误点。
  • 两版实现都已通过对拍与压力测试：显式栈保证 n=2×105n=2\times 10^5 的链不爆栈，Python 版用 CSR 数组把邻接表内存压到约 40MB。