LC 3367. 移除边之后的权重最大和

https://leetcode.cn/problems/maximize-sum-of-weights-after-edge-removals/description/

DFS框架，对于每个节点(x, parent=p)分别需要求解：

我们可以得到每个节点(a, b+w)的pair. 其中w表示(x->y)的权重。b+w表示如果可以保留的话，那么可以获得的权重。

假设我们最先假设所有的边都需要cut掉，然后再考了将每个边都加回来。每条边加回来获得的增量是(b+w-a), 所以按照这个做排序就好。 

class Solution:
    def maximizeSumOfWeights(self, edges: List[List[int]], k: int) -> int:
        n = len(edges) + 1
        adj = [[] for _ in range(n)]
        for u, v, w in edges:
            adj[u].append((v, w))
            adj[v].append((u, w))

        def dfs(x, p) -> Tuple[int, int]:
            # [a, b]
            # a: 如果必须需要cut edge(p-x)才可以满足的最大值
            # b: 如果不必须需要cut edge(p-x)就可以满足的最大值
            if len(adj[x]) == 1 and adj[x][0] == p:
                return 0, 0

            child = []
            for y, w in adj[x]:
                if y == p: continue
                a, b = dfs(y, x)
                child.append((a, b + w))

            # 不保留w, 可以得到a
            # 保留边w, 可以得到b+w
            child.sort(key=lambda x: x[1] - x[0], reverse=True)
            base = sum([x[0] for x in child])
            values = [base]
            for a, bw in child:
                base += (bw - a)
                values.append(max(values[-1], base))

            if len(values) >= (k + 1):
                return values[k], values[k - 1]
            else:
                return values[-1], values[-1]

        a, b = dfs(0, -1)
        return max(a, b)