3544. Subtree Inversion Sum LeetCode Solution
In this guide, you will get 3544. Subtree Inversion Sum LeetCode Solution with the best time and space complexity. The solution to Subtree Inversion Sum problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.
Table of Contents
- Problem Statement
- Complexity Analysis
- Subtree Inversion Sum solution in C++
- Subtree Inversion Sum solution in Java
- Subtree Inversion Sum solution in Python
- Additional Resources
Problem Statement of Subtree Inversion Sum
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] = [ui, vi] indicates an edge between nodes ui and vi.
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 * 104
edges.length == n – 1
edges[i] = [ui, vi]
0 <= ui, vi < n
nums.length == n
-5 * 104 <= nums[i] <= 5 * 104
1 <= k <= 50
The input is generated such that edges represents a valid tree.
Complexity Analysis
- Time Complexity: O(nk)
- Space Complexity: O(nk)
3544. Subtree Inversion Sum LeetCode Solution in C++
class Solution {
public:
long long subtreeInversionSum(vector<vector<int>>& edges, vector<int>& nums,
int k) {
const int n = edges.size() + 1;
vector<int> parent(n, -1);
vector<vector<int>> graph(n);
for (const vector<int>& edge : edges) {
const int u = edge[0];
const int v = edge[1];
graph[u].push_back(v);
graph[v].push_back(u);
}
vector<vector<vector<long>>> mem(
n, vector<vector<long>>(k + 1, vector<long>(2, -1)));
return dfs(graph, /*u=*/0, /*stepsSinceInversion=*/k,
/*inverted=*/false, nums, k, parent, mem);
}
private:
// Returns the maximum sum for subtree rooted at u, with `stepsSinceInversion`
// steps of inversion and `inverted` is true if the subtree is inverted.
long dfs(const vector<vector<int>>& graph, int u, int stepsSinceInversion,
bool inverted, const vector<int>& nums, int k, vector<int>& parent,
vector<vector<vector<long>>>& mem) {
if (mem[u][stepsSinceInversion][inverted] != -1)
return mem[u][stepsSinceInversion][inverted];
long num = inverted ? -nums[u] : nums[u];
long negNum = -num;
for (const int v : graph[u]) {
if (v == parent[u])
continue;
parent[v] = u;
num += dfs(graph, v, min(k, stepsSinceInversion + 1), inverted, nums, k,
parent, mem);
if (stepsSinceInversion == k)
negNum += dfs(graph, v, /*stepsSinceInversion=*/1, !inverted, nums, k,
parent, mem);
}
return mem[u][stepsSinceInversion][inverted] =
(stepsSinceInversion == k) ? max(num, negNum) : num;
}
};
/* code provided by PROGIEZ */3544. Subtree Inversion Sum LeetCode Solution in Java
class Solution {
public long subtreeInversionSum(int[][] edges, int[] nums, int k) {
final int n = edges.length + 1;
int[] parent = new int[n];
List<Integer>[] graph = new List[n];
Arrays.fill(parent, -1);
for (int i = 0; i < n; ++i)
graph[i] = new ArrayList<>();
for (int[] edge : edges) {
final int u = edge[0];
final int v = edge[1];
graph[u].add(v);
graph[v].add(u);
}
return dfs(graph, /*u=*/0, /*stepsSinceInversion=*/k,
/*inverted=*/false, nums, k, parent, new Long[n][k + 1][2]);
}
private long dfs(List<Integer>[] graph, int u, int stepsSinceInversion, boolean inverted,
int[] nums, int k, int[] parent, Long[][][] mem) {
if (mem[u][stepsSinceInversion][inverted ? 1 : 0] != null)
return mem[u][stepsSinceInversion][inverted ? 1 : 0];
long num = inverted ? -nums[u] : nums[u];
long negNum = -num;
for (final int v : graph[u]) {
if (v == parent[u])
continue;
parent[v] = u;
num += dfs(graph, v, Math.min(k, stepsSinceInversion + 1), inverted, nums, k, parent, mem);
if (stepsSinceInversion == k)
negNum += dfs(graph, v, 1, !inverted, nums, k, parent, mem);
}
return mem[u][stepsSinceInversion][inverted ? 1 : 0] =
(stepsSinceInversion == k) ? Math.max(num, negNum) : num;
}
}
// code provided by PROGIEZ3544. Subtree Inversion Sum LeetCode Solution in Python
class Solution:
def subtreeInversionSum(
self,
edges: list[list[int]],
nums: list[int],
k: int
) -> int:
n = len(edges) + 1
parent = [-1] * n
graph = [[] for _ in range(n)]
for u, v in edges:
graph[u].append(v)
graph[v].append(u)
@functools.lru_cache(None)
def dp(u: int, stepsSinceInversion: int, inverted: bool) -> int:
"""
Returns the maximum sum for subtree rooted at u, with
`stepsSinceInversion` steps of inversion and `inverted` is true if the
subtree is inverted.
"""
num = -nums[u] if inverted else nums[u]
negNum = -num
for v in graph[u]:
if v == parent[u]:
continue
parent[v] = u
num += dp(v, min(k, stepsSinceInversion + 1), inverted)
if stepsSinceInversion == k:
negNum += dp(v, 1, not inverted)
return max(num, negNum) if stepsSinceInversion == k else num
return dp(0, k, False)
# code by PROGIEZAdditional Resources
- Explore all LeetCode problem solutions at Progiez here
- Explore all problems on LeetCode website here
Happy Coding! Keep following PROGIEZ for more updates and solutions.