Hard

题目描述

给你一个以节点 0 为根的无向树，共有 n 个节点，编号从 0 到 n - 1。树由长度为 n - 1 的二维整数数组 edges 表示，其中 edges[i] = [ui, vi] 表示节点 ui 和 vi 之间有一条边。

同时给你一个长度为 n 的整数数组 nums，其中 nums[i] 表示节点 i 的值，以及一个整数 k。

你可以对节点子集执行反转操作，需要遵循以下规则：

子树反转操作：

  • 当你反转一个节点时，以该节点为根的子树中的每个值都会乘以 -1。

反转的距离约束：

  • 只有当一个节点与任何其他已反转节点"足够远"时，你才能反转该节点。
  • 具体而言，如果你反转两个节点 a 和 b，使得其中一个是另一个的祖先（即 LCA(a, b) = a 或 LCA(a, b) = b），那么它们之间的距离（它们之间唯一路径上的边数）必须至少为 k。

返回应用反转操作后树的节点值的最大可能和。

示例 1：

输入：edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], nums = [4,-8,-6,3,7,-2,5], k = 2
输出：27
解释：在节点 0、3、4 和 6 处应用反转操作。
最终 nums 数组为 [-4, 8, 6, 3, 7, 2, 5]，总和为 27。

示例 2：

输入：edges = [[0,1],[1,2],[2,3],[3,4]], nums = [-1,3,-2,4,-5], k = 2
输出：9
解释：在节点 4 处应用反转操作。
最终 nums 数组变为 [-1, 3, -2, 4, 5]，总和为 9。

示例 3：

输入：edges = [[0,1],[0,2]], nums = [0,-1,-2], k = 3
输出：3
解释：在节点 1 和 2 处应用反转操作。

约束条件：

  • 2 <= n <= 5 * 10^4
  • edges.length == n - 1
  • edges[i] = [ui, vi]
  • 0 <= ui, vi < n
  • nums.length == n
  • -5 * 10^4 <= nums[i] <= 5 * 10^4
  • 1 <= k <= 50
  • 输入保证 edges 表示一棵有效的树。

解题思路

这是一个基于树的动态规划问题。我们需要考虑在满足距离约束的情况下，如何选择反转节点来最大化总和。

核心思路：

  1. 状态定义：使用三维 DP 状态 dp[node][parity][dist]，其中：

    • node：当前节点
    • parity：祖先反转的奇偶性（0表示偶数次，1表示奇数次）
    • dist：从当前节点到最近祖先反转节点的距离

  2. 转移方程：对于每个节点，我们有两种选择：

    • 不反转当前节点：继承父节点的状态，距离加1
    • 反转当前节点：需要满足距离约束（dist >= k），重置距离为0，奇偶性翻转

  3. 值计算：根据奇偶性决定当前子树的符号，如果奇数次反转则取负值，偶数次则保持原值。

  4. 边界条件：当距离不足 k 时，不能进行反转操作。

算法步骤：

  1. 构建邻接表表示树
  2. 从根节点开始 DFS，维护三维 DP 状态
  3. 对每个节点尝试反转和不反转两种情况
  4. 递归处理子节点，累加最优结果
  5. 返回根节点的最优解

时间复杂度主要由状态数量决定，每个节点最多有 2×k 个状态。

代码实现

class Solution {
public:
    long long subtreeInversionSum(vector<vector<int>>& edges, vector<int>& nums, int k) {
        int n = nums.size();
        vector<vector<int>> adj(n);
        
        for (auto& edge : edges) {
            adj[edge[0]].push_back(edge[1]);
            adj[edge[1]].push_back(edge[0]);
        }
        
        // dp[node][parity][dist] = maximum sum
        vector<vector<vector<long long>>> memo(n, vector<vector<long long>>(2, vector<long long>(k + 1, LLONG_MIN)));
        
        function<long long(int, int, int, int)> dfs = [&](int node, int parent, int parity, int dist) -> long long {
            if (memo[node][parity][dist] != LLONG_MIN) {
                return memo[node][parity][dist];
            }
            
            long long result = LLONG_MIN;
            
            // Option 1: Don't invert current node
            long long sum1 = parity ? -nums[node] : nums[node];
            for (int child : adj[node]) {
                if (child != parent) {
                    int newDist = min(dist + 1, k);
                    sum1 += dfs(child, node, parity, newDist);
                }
            }
            result = max(result, sum1);
            
            // Option 2: Invert current node (if distance constraint is satisfied)
            if (dist >= k) {
                long long sum2 = parity ? nums[node] : -nums[node];
                for (int child : adj[node]) {
                    if (child != parent) {
                        sum2 += dfs(child, node, 1 - parity, 0);
                    }
                }
                result = max(result, sum2);
            }
            
            return memo[node][parity][dist] = result;
        };
        
        return dfs(0, -1, 0, k);
    }
};

class Solution:
    def subtreeInversionSum(self, edges: List[List[int]], nums: List[int], k: int) -> int:
        n = len(nums)
        adj = [[] for _ in range(n)]
        
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        
        memo = {}
        
        def dfs(node, parent, parity, dist):
            if (node, parity, dist) in memo:
                return memo[(node, parity, dist)]
            
            result = float('-inf')
            
            # Option 1: Don't invert current node
            sum1 = -nums[node] if parity else nums[node]
            for child in adj[node]:
                if child != parent:
                    new_dist = min(dist + 1, k)
                    sum1 += dfs(child, node, parity, new_dist)
            result = max(result, sum1)
            
            # Option 2: Invert current node (if distance constraint is satisfied)
            if dist >= k:
                sum2 = nums[node] if parity else -nums[node]
                for child in adj[node]:
                    if child != parent:
                        sum2 += dfs(child, node, 1 - parity, 0)
                result = max(result, sum2)
            
            memo[(node, parity, dist)] = result
            return result
        
        return dfs(0, -1, 0, k)

public class Solution {
    private Dictionary<(int, int, int), long> memo;
    private List<List<int>> adj;
    private int[] nums;
    private int k;
    
    public long SubtreeInversionSum(int[][] edges, int[] nums, int k) {
        int n = nums.Length;
        this.nums = nums;
        this.k = k;
        this.memo = new Dictionary<(int, int, int), long>();
        this.adj = new List<List<int>>();
        
        for (int i = 0; i < n; i++) {
            adj.Add(new List<int>());
        }
        
        foreach (var edge in edges) {
            adj[edge[0]].Add(edge[1]);
            adj[edge[1]].Add(edge[0]);
        }
        
        return DFS(0, -1, 0, k);
    }
    
    private long DFS(int node, int parent, int parity, int dist) {
        var key = (node, parity, dist);
        if (memo.ContainsKey(key)) {
            return memo[key];
        }
        
        long result = long.MinValue;
        
        // Option 1: Don't invert current node
        long sum1 = parity == 1 ? -nums[node] : nums[node];
        foreach (int child in adj[node]) {
            if (child != parent) {
                int newDist = Math.Min(dist + 1, k);
                sum1 += DFS(child, node, parity, newDist);
            }
        }
        result = Math.Max(result, sum1);
        
        // Option 2: Invert current node (if distance constraint is satisfied)
        if (dist >= k) {
            long sum2 = parity == 1 ? nums[node] : -nums[node];
            foreach (int child in adj[node]) {
                if (child != parent) {
                    sum2 += DFS(child, node, 1 - parity, 0);
                }
            }
            result = Math.Max(result, sum2);
        }
        
        memo[key] = result;
        return result;
    }
}

var subtreeInversionSum = function(edges, nums, k) {
    const n = nums.length;
    const adj = Array.from({length: n}, () => []);
    
    for (const [u, v] of edges) {
        adj[u].push(v);
        adj[v].push(u);
    }
    
    const memo = new Map();
    
    function dfs(node, parent, parity, dist) {
        const key = `${node},${parity},${dist}`;
        if (memo.has(key)) {
            return memo.get(key);
        }
        
        let result = -Infinity;
        
        // Option 1: Don't invert current node
        let sum1 = parity ? -nums[node] : nums[node];
        for (const child of adj[node]) {
            if (child !== parent) {
                const newDist = Math.min(dist + 1, k);
                sum1 += dfs(child, node, parity, newDist);
            }
        }
        result = Math.max(result, sum1);
        
        // Option 2: Invert current node (if distance constraint is satisfied)
        if (dist >= k) {
            let sum2 = parity ? nums[node] : -nums[node];
            for (const child of adj[node]) {
                if (child !== parent) {
                    sum2 += dfs(child, node, 1 - parity, 0);
                }
            }
            result = Math.max(result, sum2);
        }
        
        memo.set(key, result);
        return result;
    }
    
    return dfs(0, -1, 0, k);
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型分析
时间复杂度O(n × k)
空间复杂度O(n × k)

说明：

  • 时间复杂度：每个节点最多有 2×k 个状态，每个状态计算一次，总共 O(n × k) 个状态
  • 空间复杂度：记忆化存储需要 O(n × k) 的空间，递归栈深度最多为 O(n)