Subtree Inversion Sum

Hard

#3155Time: O(2^N) in the worst case. For each node, there are up to two branches of recursion. In a path-like tree, this can lead to exploring a number of paths exponential in the number of nodes `N`.Space: O(N) for the recursion call stack depth in a skewed tree.

Patterns

Dynamic Programming

Algorithms

Depth-First Search

Data structures

Array Tree

Prompt

You are given an undirected tree rooted at node 0, with n nodes numbered from 0 to n - 1. The tree is represented by a 2D integer array edges of length n - 1, where edges[i] = [u_i, v_i] indicates an edge between nodes u_i and v_i.

You are also given an integer array nums of length n, where nums[i] represents the value at node i, and an integer k.

You may perform inversion operations on a subset of nodes subject to the following rules:

Subtree Inversion Operation:
- When you invert a node, every value in the subtree rooted at that node is multiplied by -1.
Distance Constraint on Inversions:
- You may only invert a node if it is "sufficiently far" from any other inverted node.
- Specifically, if you invert two nodes a and b such that one is an ancestor of the other (i.e., if LCA(a, b) = a or LCA(a, b) = b), then the distance (the number of edges on the unique path between them) must be at least k.

Return the maximum possible sum of the tree's node values after applying inversion operations.

Example 1:

Input: edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], nums = [4,-8,-6,3,7,-2,5], k = 2

Output: 27

Explanation:

Apply inversion operations at nodes 0, 3, 4 and 6.
The final nums array is [-4, 8, 6, 3, 7, 2, 5], and the total sum is 27.

Example 2:

Input: edges = [[0,1],[1,2],[2,3],[3,4]], nums = [-1,3,-2,4,-5], k = 2

Output: 9

Explanation:

Apply the inversion operation at node 4.
The final nums array becomes [-1, 3, -2, 4, 5], and the total sum is 9.

Example 3:

Input: edges = [[0,1],[0,2]], nums = [0,-1,-2], k = 3

Output: 3

Explanation:

Apply inversion operations at nodes 1 and 2.

Constraints:

2 <= n <= 5 * 10⁴
edges.length == n - 1
edges[i] = [u_i, v_i]
0 <= u_i, v_i < n
nums.length == n
-5 * 10⁴ <= nums[i] <= 5 * 10⁴
1 <= k <= 50
The input is generated such that edges represents a valid tree.

Approaches

2 approaches with complexity analysis and trade-offs.

A straightforward approach is to use recursion to explore all possible valid sets of node inversions. This method, also known as backtracking, systematically tries every choice (invert or not invert) for each node and, after exploring all nodes, calculates the resulting sum. It keeps track of the maximum sum found.

Algorithm

Define a recursive function solve(u, p, d).
Base case: If u is a leaf node, calculate its value based on the two choices (invert or not) and return the maximum.
Recursive step for an internal node u:

Calculate the sum for the 'don't invert u' case: value_at_u + sum(solve(child, p, d+1)) for all children.
If d >= k, calculate the sum for the 'invert u' case: inverted_value_at_u + sum(solve(child, (p+1)%2, 1)) for all children.
Return the maximum of the calculated sums.

The initial call is solve(0, 0, k).

Walkthrough

We can define a recursive function, say solve(u, p, d), that computes the maximum possible sum for the subtree rooted at u. The state for the recursion needs to carry information from the ancestors that affects decisions within the current subtree.

u: The current node being processed.
p: The parity (0 for even, 1 for odd) of inversion operations applied on the path from the root to the parent of u.
d: The distance to the nearest ancestor of u that has been inverted. We can use a special value (e.g., k) to signify that no ancestor has been inverted or the nearest one is at a distance of k or more.

For each node u, we explore two possibilities:

Do not invert u: The node's value contributes nums[u] * (-1)^p to the sum. We then recursively call the function for all children of u, passing down the same parity p and an incremented distance d+1.
Invert u: This choice is only available if the distance constraint is met, i.e., d >= k. If we invert u, its value contributes nums[u] * (-1)^(p+1). For its children, the new parity becomes (p+1)%2, and the distance to the nearest inverted ancestor resets to 1.

The function returns the maximum sum obtained from the valid choices. The recursion starts from the root solve(0, 0, k).

Without memoization, this approach recomputes the solution for the same subproblems (same node u with same p and d) multiple times, leading to an exponential time complexity.

Complexity

Time

O(2^N) in the worst case. For each node, there are up to two branches of recursion. In a path-like tree, this can lead to exploring a number of paths exponential in the number of nodes `N`.

Space

O(N) for the recursion call stack depth in a skewed tree.

Trade-offs

Pros

Conceptually simple and directly follows the problem definition.

Cons

Highly inefficient and will time out for the given constraints due to redundant computations.

This approach optimizes the backtracking solution by using memoization to store and reuse the results of subproblems. The state of a subproblem is defined by the current node, the inversion parity from its ancestors, and the distance to the nearest inverted ancestor. This avoids redundant calculations and reduces the complexity significantly.

Algorithm

First, convert the given edges into a directed tree representation (e.g., an adjacency list where adj[u] contains children of u), rooted at 0.
Initialize a 3D memoization table memo[n][2][k+1] with a sentinel value to indicate that a state has not been computed.
Define a recursive function solve(u, p, d) that computes the maximum sum for the subtree at u.
In solve(u, p, d), first check if memo[u][p][d] already contains a computed result. If so, return it.
Calculate the result for the 'don't invert u' case by summing nums[u]'s contribution and the results of recursive calls for its children: solve(v, p, min(d+1, k)).
If d >= k, calculate the result for the 'invert u' case similarly: solve(v, 1-p, 1).
Store the maximum of the possible outcomes in memo[u][p][d] and return it.
The final answer is obtained by calling solve(0, 0, k).

Walkthrough

The problem has optimal substructure and overlapping subproblems, making it suitable for dynamic programming. We define a DP state dp[u][p][d] which stores the maximum sum of the subtree rooted at u under specific ancestor conditions.

u: The current node index (0 to n-1).
p: The parity of inversions from ancestors (0 or 1).
d: The distance to the nearest inverted ancestor (1 to k). d=k is a special state indicating the distance is at least k or no ancestor is inverted.

We use a post-order traversal (DFS) to compute the DP table from leaves up to the root. For each state (u, p, d), we calculate the maximum possible sum by considering two choices for node u:

Don't invert u: The total sum is (nums[u] * (-1)^p) + sum(dp[v][p][min(d+1, k)]) for all children v of u.
Invert u: This is only possible if d >= k. The total sum is (nums[u] * (-1)^(p+1)) + sum(dp[v][(p+1)%2][1]) for all children v of u.

The value dp[u][p][d] is the maximum of these choices. The final answer is the result of the initial call for the root node, dp[0][0][k].

1class Solution {2    private List<List<Integer>> adj;3    private long[][][] memo;4    private int[] nums;5    private int k;6 7    public long subtreeInversionSum(int[][] edges, int[] nums, int k) {8        int n = nums.length;9        this.adj = new ArrayList<>();10        for (int i = 0; i < n; i++) {11            adj.add(new ArrayList<>());12        }13        // Build a directed tree (parent->child) from the undirected edges, rooted at 014        List<List<Integer>> undirectedAdj = new ArrayList<>();15        for (int i = 0; i < n; i++) undirectedAdj.add(new ArrayList<>());16        for (int[] edge : edges) {17            undirectedAdj.get(edge[0]).add(edge[1]);18            undirectedAdj.get(edge[1]).add(edge[0]);19        }20 21        boolean[] visited = new boolean[n];22        Queue<Integer> q = new LinkedList<>();23        q.add(0);24        visited[0] = true;25        while (!q.isEmpty()) {26            int u = q.poll();27            for (int v : undirectedAdj.get(u)) {28                if (!visited[v]) {29                    visited[v] = true;30                    adj.get(u).add(v);31                    q.add(v);32                }33            }34        }35 36        this.nums = nums;37        this.k = k;38        this.memo = new long[n][2][k + 1];39        for (long[][] plane : memo) {40            for (long[] row : plane) {41                Arrays.fill(row, Long.MIN_VALUE);42            }43        }44 45        return solve(0, 0, k);46    }47 48    private long solve(int u, int p, int d) {49        if (memo[u][p][d] != Long.MIN_VALUE) {50            return memo[u][p][d];51        }52 53        // Option 1: Don't invert node u54        long childrenSumDontInvert = 0;55        for (int v : adj.get(u)) {56            childrenSumDontInvert += solve(v, p, Math.min(d + 1, k));57        }58        long valAtU_dontInvert = (p % 2 == 0) ? nums[u] : -nums[u];59        long totalSumDontInvert = valAtU_dontInvert + childrenSumDontInvert;60 61        // Option 2: Invert node u62        long totalSumInvert = Long.MIN_VALUE;63        if (d >= k) {64            long childrenSumInvert = 0;65            for (int v : adj.get(u)) {66                childrenSumInvert += solve(v, 1 - p, 1);67            }68            long valAtU_invert = ((p + 1) % 2 == 0) ? nums[u] : -nums[u];69            totalSumInvert = valAtU_invert + childrenSumInvert;70        }71 72        return memo[u][p][d] = Math.max(totalSumDontInvert, totalSumInvert);73    }74}

Complexity

Time

O(N * K). Each state `(u, p, d)` is computed once. There are `N * 2 * (K+1)` such states. The work for each state is proportional to the number of its children. The total time across all states is `sum_{u,p,d} (1 + degree(u))`, which is `O(N*K)` because `sum(degree(u)) = N-1`.

Space

O(N * K) for the memoization table. An additional `O(N)` is used for the adjacency list and the recursion stack.

Trade-offs

Pros

Efficient and guaranteed to find the optimal solution.
Handles the constraints of the problem effectively.

Cons

Requires more memory due to the 3D DP table.
Slightly more complex to implement than the naive backtracking approach.

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem