Hard

题目描述

给定一个有 n 个节点的无向树,节点编号从 0n-1。树用长度为 n-1 的二维整数数组 edges 表示,其中 edges[i] = [ai, bi] 表示节点 aibi 之间有一条边。

同时给定一个长度为 n 的整数数组 good,其中 good[i]1 表示第 i 个节点是好节点,为 0 表示是坏节点。

定义子图的分数为该子图中好节点的数量减去坏节点的数量。

对于每个节点 i,找出所有包含节点 i 的连通子图中的最大可能分数。

返回一个长度为 n 的整数数组,其中第 i 个元素是节点 i 的最大分数。

子图是原图的顶点和边的子集构成的图。

连通子图是指其中任意两个顶点之间都可以通过该子图的边相互到达的子图。

示例 1:

输入:n = 3, edges = [[0,1],[1,2]], good = [1,0,1]

输出:[1,1,1]

示例 2:

输入:n = 5, edges = [[1,0],[1,2],[1,3],[3,4]], good = [0,1,0,1,1]

输出:[2,3,2,3,3]

示例 3:

输入:n = 2, edges = [[0,1]], good = [0,0]

输出:[-1,-1]

约束条件:

  • 2 <= n <= 10^5
  • edges.length == n - 1
  • edges[i] = [ai, bi]
  • 0 <= ai, bi < n
  • good.length == n
  • 0 <= good[i] <= 1
  • 输入保证 edges 表示一个有效的树

解题思路

这是一道典型的树形动态规划(Tree DP)+ 换根动态规划(Rerooting DP)问题。

核心思路是:对于每个节点,我们需要考虑所有包含该节点的连通子图,找出分数最大的那个。

主要思路:

  1. 第一次DFS(向下收集):以任意节点为根,计算每个节点作为子树根时,从该节点向下能得到的最大分数。对于每个节点,我们需要考虑:

    • 只包含当前节点的子图
    • 包含当前节点及其某些子树的子图
  2. 第二次DFS(换根传递):利用换根技巧,将父节点的信息传递给子节点。当我们从父节点转移到子节点时,需要重新计算包含"父节点方向"的最优解。

关键观察:

  • 每个节点的贡献值是 good[i] == 1 ? 1 : -1
  • 对于任何子树,如果其总贡献为负,我们可以选择不包含它
  • 我们需要维护每个节点向下的最大分数,以及向上(通过父节点)的最大分数

算法步骤:

  1. 构建邻接表表示树
  2. 第一次DFS:计算每个节点向下的最大分数
  3. 第二次DFS:利用换根技巧计算最终答案

时间复杂度为O(n),空间复杂度为O(n)。

代码实现

class Solution {
public:
    vector<int> maxSubgraphScore(int n, vector<vector<int>>& edges, vector<int>& good) {
        vector<vector<int>> adj(n);
        for (auto& edge : edges) {
            adj[edge[0]].push_back(edge[1]);
            adj[edge[1]].push_back(edge[0]);
        }
        
        vector<int> down(n), ans(n);
        
        function<void(int, int)> dfs1 = [&](int u, int p) {
            down[u] = good[u] == 1 ? 1 : -1;
            for (int v : adj[u]) {
                if (v != p) {
                    dfs1(v, u);
                    down[u] = max(down[u], down[u] + down[v]);
                }
            }
        };
        
        function<void(int, int, int)> dfs2 = [&](int u, int p, int up) {
            ans[u] = max(down[u], up);
            
            vector<int> prefix(adj[u].size() + 1, 0);
            vector<int> suffix(adj[u].size() + 1, 0);
            vector<int> children;
            
            for (int v : adj[u]) {
                if (v != p) {
                    children.push_back(v);
                }
            }
            
            int base = good[u] == 1 ? 1 : -1;
            
            for (int i = 0; i < children.size(); i++) {
                prefix[i + 1] = max(prefix[i], prefix[i] + down[children[i]]);
            }
            
            for (int i = children.size() - 1; i >= 0; i--) {
                suffix[i] = max(suffix[i + 1], suffix[i + 1] + down[children[i]]);
            }
            
            for (int i = 0; i < children.size(); i++) {
                int new_up = base;
                if (up > 0) new_up = max(new_up, new_up + up);
                if (prefix[i] > 0) new_up = max(new_up, new_up + prefix[i]);
                if (suffix[i + 1] > 0) new_up = max(new_up, new_up + suffix[i + 1]);
                
                dfs2(children[i], u, new_up);
            }
        };
        
        dfs1(0, -1);
        dfs2(0, -1, good[0] == 1 ? 1 : -1);
        
        return ans;
    }
};
class Solution:
    def maxSubgraphScore(self, n: int, edges: List[List[int]], good: List[int]) -> List[int]:
        adj = [[] for _ in range(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        
        down = [0] * n
        ans = [0] * n
        
        def dfs1(u, p):
            down[u] = 1 if good[u] == 1 else -1
            for v in adj[u]:
                if v != p:
                    dfs1(v, u)
                    down[u] = max(down[u], down[u] + down[v])
        
        def dfs2(u, p, up):
            ans[u] = max(down[u], up)
            
            children = [v for v in adj[u] if v != p]
            if not children:
                return
            
            base = 1 if good[u] == 1 else -1
            
            # 计算前缀和后缀的最大值
            prefix = [0] * (len(children) + 1)
            suffix = [0] * (len(children) + 1)
            
            for i in range(len(children)):
                prefix[i + 1] = max(prefix[i], prefix[i] + down[children[i]])
            
            for i in range(len(children) - 1, -1, -1):
                suffix[i] = max(suffix[i + 1], suffix[i + 1] + down[children[i]])
            
            for i, child in enumerate(children):
                new_up = base
                if up > 0:
                    new_up = max(new_up, new_up + up)
                if prefix[i] > 0:
                    new_up = max(new_up, new_up + prefix[i])
                if suffix[i + 1] > 0:
                    new_up = max(new_up, new_up + suffix[i + 1])
                
                dfs2(child, u, new_up)
        
        dfs1(0, -1)
        dfs2(0, -1, 1 if good[0] == 1 else -1)
        
        return ans
public class Solution {
    public int[] MaxSubgraphScore(int n, int[][] edges, int[] good) {
        var adj = new List<int>[n];
        for (int i = 0; i < n; i++) {
            adj[i] = new List<int>();
        }
        
        foreach (var edge in edges) {
            adj[edge[0]].Add(edge[1]);
            adj[edge[1]].Add(edge[0]);
        }
        
        var down = new int[n];
        var ans = new int[n];
        
        void Dfs1(int u, int p) {
            down[u] = good[u] == 1 ? 1 : -1;
            foreach (int v in adj[u]) {
                if (v != p) {
                    Dfs1(v, u);
                    down[u] = Math.Max(down[u], down[u] + down[v]);
                }
            }
        }
        
        void Dfs2(int u, int p, int up) {
            ans[u] = Math.Max(down[u], up);
            
            var children = new List<int>();
            foreach (int v in adj[u]) {
                if (v != p) {
                    children.Add(v);
                }
            }
            
            if (children.Count == 0) return;
            
            int baseVal = good[u] == 1 ? 1 : -1;
            
            var prefix = new int[children.Count + 1];
            var suffix = new int[children.Count + 1];
            
            for (int i = 0; i < children.Count; i++) {
                prefix[i + 1] = Math.Max(prefix[i], prefix[i] + down[children[i]]);
            }
            
            for (int i = children.Count - 1; i >= 0; i--) {
                suffix[i] = Math.Max(suffix[i + 1], suffix[i + 1] + down[children[i]]);
            }
            
            for (int i = 0; i < children.Count; i++) {
                int newUp = baseVal;
                if (up > 0) newUp = Math.Max(newUp, newUp + up);
                if (prefix[i] > 0) newUp = Math.Max(newUp, newUp + prefix[i]);
                if (suffix[i + 1] > 0) newUp = Math.Max(newUp, newUp + suffix[i + 1]);
                
                Dfs2(children[i], u, newUp);
            }
        }
        
        Dfs1(0, -1);
        Dfs2(0, -1, good[0] == 1 ? 1 : -1);
        
        return ans;
    }
}
var maxSubgraphScore = function(n, edges, good) {
    const adj = Array.from({length: n}, () => []);
    for (const [u, v] of edges) {
        adj[u].push(v);
        adj[v].push(u);
    }
    
    const down = new Array(n).fill(0);
    const ans = new Array(n).fill(0);
    
    function dfs1(u, p) {
        down[u] = good[u]

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n)每个节点被访问常数次,每条边被遍历常数次
空间复杂度O(n)递归栈深度最大为O(n),额外数组空间为O(n)