从大厂员工吐槽到算法实战：均匀树划分的DFS解法与避坑指南

一位大厂哥们在网上开麦：“我人生很失败。”理由也够扎心：赚到一千多万去买房，结果房价一拐弯，账面先给他上了一课；娃的成绩稳坐班级倒数，辅导作业像线上故障排查，越查越多红；媳妇天天抱怨，他夹在中间跟夹在两套KPI里一样喘不过气。

我觉得这事最狠的地方在于，他明明很努力，却被房价、分数、情绪三连击。

好在这种技术社区有个优点：你发出来，会发现自己不是唯一一个被生活追着打的，顺便还能捡到几条真能用的解法。这不，连带着他这吐槽，我脑子里蹦出来的不是别的，而是之前啃过的一道算法面试题——均匀树划分。

面试题：均匀树划分

我跟你说啊，这个“均匀树划分”这玩意儿，第一次看到题目我还以为是领导让我把部门人员平均拆成几个小组，结果一看输入输出：一棵树、每个点一个权值、问你最多能把树砍成多少块，每块权值和都一样。嗯…还是算法世界的事，比拆团队安全点，至少不会被 HR 请去喝茶。

大概模型你脑子里过一下哈：

• N 个节点，一棵树，边都是无向的
• 每个节点有个值 a[i]
• 你可以删一些边，把这棵树拆成若干连通块
• 要求每个连通块里所有点的权值和都一样
• 让你求：最多能划成多少块

这个地方最容易脑抽的地方是：你要先猜那个“每块的和”是多少，对吧？你总不能一个一个尝，从 1 到 totalSum 枚举，那人都没了。

我的思路当时是这么来的：有一次写 Postgres 和 MySQL 压测的时候，不也是先看总 QPS，然后再想着怎么切分业务嘛，宏观总量先算清楚，后面才好细化。这题也一样，先搞一个 totalSum，把所有点加起来。

然后你想啊，如果最后能划成 k 块，每块的和就是 target = totalSum / k。那 totalSum 至少要能被 k 整除吧？不然那就是纯幻想。

比较暴力但好写的办法是：

• 从 k = N 一直往下试到 1
• 只有当 totalSum % k == 0 的时候，我们才试一下能不能把树切成 k 块
• 每次 DFS 一遍树，看有多少子树的和刚好等于 target = totalSum / k
• 如果刚好数出来有 k 个这样的子树，那就说明可以切成 k 块，答案就是删 k - 1 条边

为啥是 k-1？你想象每多一块就要多砍一条边，原来是 1 块不砍，2 块砍 1 条，3 块砍 2 条，很自然嘛。

我用 Java 写的时候，大概是下面这样的，代码你直接看，先有个整体直觉就行，细节别纠结太多：

// 均匀树划分，小心一点就行，不难
public class EqualTreeCut {
    int n;
    List<Integer>[] g;
    int[] val;
    long total;
    long target;
    int pieces; // 数出来有多少块

    long dfs(int u, int parent) {
        long sum = val[u];
        for (int v : g[u]) {
            if (v == parent) continue;
            long childSum = dfs(v, u);
            sum += childSum;
        }
        // 如果这个子树刚好等于 target，我们就当它是一块，往上返回 0
        if (sum == target) {
            pieces++;
            return 0;
        }
        return sum;
    }

    // 主逻辑，返回最多能切成几块
    public int solve() {
        total = 0;
        for (int i = 1; i <= n; i++) {
            total += val[i];
        }
        // 从大到小试 k，先试能不能分成很多块
        for (int k = n; k >= 1; k--) {
            if (total % k != 0) continue;
            target = total / k;
            pieces = 0;
            dfs(1, 0);
            if (pieces == k) {
                return k;
            }
        }
        return 1;
    }
}

你看这个 dfs 里面那个小心机：

• 先把孩子的和都加上来
• 如果当前子树的 sum 刚好等于 target，就把 pieces++ 然后返回 0，意思是：这块我就自己成团了，往上就别再加我了
• 如果没到 target，就把 sum 往上扔，让父节点继续攒

这个写法有两个坑，我自己踩过：

第一个是数据范围。题目一般喜欢搞点 1e5、1e6 那种，如果点权再是 1e9 级别，你 int 直接就爆了。像上面这样我直接全用 long，图省心，别问为啥，问就是我因为这个多查过一次线上 SQL 日志，怀疑人生。

第二个是递归深度。树要是特别像一条链，N 比较大，Java 默认栈深度可能会喊“我不行了”。这个一般面试题里不会搞那么狠，真要上生产我会乖乖改成手写栈：

// 伪代码示意，真写起来略啰嗦
Deque<Integer> stack = new ArrayDeque<>();
stack.push(root);
// 先搞个后序遍历顺序 list，再反着跑一遍累加子树和

还有一个点，挺多同学一开始会误会，以为只要某些子树和是 totalSum 的因子就行，但这里有个“连通块”的限制，你不能随便拿几个点凑个数，必须是一棵棵子树，所以 DFS 这种自底向上的累加特别合适。

我当时第一次做这个题的时候，脑子一热还写了个版本，枚举 target 然后在 DFS 里顺便剪枝，结果写着写着发现：诶这不就是现在这个版本嘛，只不过换了个说法而已…你写算法时间久了就会发现，90% 的“巧妙思路”，其实就是“多绕了一圈又回来了”。