HackerRank CutTree问题解法代码的工作原理问询

先明确下我们要解决的问题：

给定一棵有n个节点的树T，统计有多少个子树T'满足：连接T'和原树剩余部分（T-T'）的边数不超过K。

输入：第一行是n和K，接下来n-1行是树的边（节点编号1~n）；输出：符合条件的子树总数。

下面我会一步步拆解这段代码的工作逻辑，帮你理解它是怎么解决这个问题的：

整体核心思路 #

这段代码用**深度优先搜索（DFS）+ 动态规划（DP）**的组合思路。本质是遍历树的每个节点时，用DP字典记录「以当前节点为根的子树中，不同切割边数对应的子树数量」，然后把所有切割边数≤K的情况累加，最终得到总数。

代码逐部分解析 #

1. 树的邻接表构建 #

g = [[] for _ in range(n)]
for i in edges:
    g[i[0]-1].append(i[1]-1)
    g[i[1]-1].append(i[0]-1)

输入的节点是1开头的，这里转成0开头的邻接表g，方便用数组索引操作。每个位置存储对应节点的所有相邻节点。

2. 全局变量初始化 #

global ans
ans = 1

ans初始值设为1，是因为我们先把整棵树本身算进去了——整棵树作为子树时，和剩余部分的连接边数是0，肯定符合条件。

3. multiply函数：合并子树的DP状态 #

def multiply(x, y):
    ans = defaultdict(lambda: 0)
    for k, v in x.items():
        for k1, v1 in y.items(): ans[k+k1-1] += v*v1
    for k, v in ans.items():
        if k in x: x[k] += v
        else: x[k] = v

这个函数是用来合并两个子树的DP状态的：

• x和y都是字典，键是切割边数，值是对应切割边数下的子树数量。
• 为什么是k+k1-1？假设子树A需要切k条边才能独立，子树B需要切k1条边才能独立。当把它们合并到同一个父节点下时，原本各自要切的边里，有一条是连接到父节点的边被重复计算了，所以要减1。比如：子树A切k条边（其中一条是连父节点的），子树B切k1条边（其中一条也是连父节点的），合并后这两条边其实只需要切一次，总切割数就是k+k1-1。
• 合并后的结果会更新到x里，让x变成合并后的状态字典。

4. dfs函数：DFS遍历+DP统计 #

def dfs(i,p):
    global ans
    # 处理叶子节点（邻接列表只有父节点）
    if g[i] == [p]:
        ans += 1
        return {0:1}
    # x标记是否是根节点：根节点没有父节点，所以邻接边数要减去0（x=1时，len(g[i])-x就是去掉父节点的数量）
    x = 1 if i else 0
    # 初始状态：当前节点单独作为子树时，需要切割的边数 = 总邻接边数 - 父节点的边数（根节点没有父节点，所以减x）
    res = {len(g[i])-x : 1}
    # 遍历所有子节点（排除父节点）
    for nxt in g[i]:
        if nxt != p: 
            # 递归遍历子节点，合并子节点的DP状态到当前res
            multiply(res, dfs(nxt, i))
    # 累加当前节点所有满足「切割边数≤K」的子树数量到ans
    ans += sum(((v if k <= K else 0) for k, v in res.items()))
    # 返回当前节点的DP状态，供父节点合并使用
    return res

这段是核心逻辑，我再拆细讲：

• 叶子节点处理：如果当前节点是叶子（邻接列表只有父节点），那么这个节点单独作为子树是符合条件的，所以ans +=1。返回的{0:1}表示：当这个叶子被包含在父节点的子树中时，不需要额外切割边（父节点会处理和它的连接）。
• 初始状态设置：res初始化为{len(g[i])-x : 1}，意思是「当前节点单独作为子树时，需要切割掉除了父节点之外的所有边，这样的子树只有1个（就是节点自己）」。
• 合并子节点状态：遍历每个子节点，递归调用DFS得到子节点的DP状态，然后用multiply合并到当前res里。这样res就包含了所有「以当前节点为根，包含任意子节点组合」的子树的切割边数和对应数量。
• 累加符合条件的子树：遍历res的所有条目，把切割边数≤K的子树数量加到ans里——这些都是以当前节点为根的、满足要求的子树。
• 返回状态：把当前节点的DP状态返回给父节点，让父节点可以继续合并其他子节点的状态。

5. 启动DFS并返回结果 #

dfs(0,-1)
return ans

从0号节点（对应输入的1号节点）开始DFS遍历整棵树，最后返回累计的ans，就是所有符合条件的子树总数。

备注：内容来源于stack exchange，提问作者Irina