2382. Maximum Segment Sum After Removals

HardUnion FindArrayOrdered SetPrefix Sum

Problem Description

In this LeetCode problem, we are given two arrays: nums and removeQueries, both containing n elements. The nums array represents a sequence of integers, and the removeQueries array contains indices representing the order in which elements will be removed from the nums array. Every time an element is removed from nums, it potentially splits the array into different continuous segments. A segment is defined as a consecutive sequence of positive integers within nums, and a segment sum is the total sum of all elements in a segment.

Our task is to calculate and return an array answer of length n, where answer[i] represents the maximum segment sum after performing the i-th removal specified in removeQueries. The constraints guarantee that each index from removeQueries is unique, ensuring that no index is removed more than once.

Intuition

The solution employs a Union-Find data structure, which is a typical approach for managing elements in disjoint sets and is an excellent choice for problems that involve dynamic connectivity queries, including segment merging or splitting tasks like this one.

Given that removing an element from nums may split the array into different segments, we can think of each element as initially being its separate segment. As we process removals starting from the last one and going backward, we simulate merging adjacent segments: whenever two positive elements become adjacent (after another unrelated removal), their segments merge into one.

The find helper function is a typical Union-Find operation - it finds the representative (parent) of a segment. If a segment's parent is not itself, it means the segment is connected to another one. The find function recursively finds the root parent and employs path compression for optimization, which flattens the structure of the tree, leading to very fast subsequent find operations.

The merge function uses the find function to get the root parents of two segments and then merges them by making one the parent of the other. It also updates the sum of the merged segment to include the sum of both segments.

We initialize two arrays: p, which tracks the parent indices for the Union-Find structure, and s, which tracks the sum of elements for each segment. The sums in s are initiated to 0 and get updated when elements are 'reactivated' by being encountered in reverse order in removeQueries.

We also maintain a variable mx to keep track of the current maximum segment sum after each removal. We start with mx set to 0 because, with all elements removed, the maximum segment sum is 0.

The solution iterates through the removeQueries array in reverse, 'restoring' the removed element at each step and processing potential segment merges on its left and right (if any), and then determines the new maximum segment sum. The updated maximum is stored at the appropriate index in the resulting array ans, which is then returned at the end of the operation.

This reverse process enables us to efficiently handle the removal of elements while keeping the maximum segment sum updated, without recomputing sums for all segments from scratch after each removal.

Learn more about Union Find and Prefix Sum patterns.

Solution Approach

The solution approach leverages the Union-Find data structure which is particularly useful in problems involving connectivity and dynamic components, like segments in an array in this case. Here, a Union-Find is used to keep track of segments within nums as removals occur.

Initially, we create two arrays, p and s, of length n. The p array represents the parent of each segment for Union-Find operations, where each element is initially its own parent, indicating that they are all separate segments. The s array represents the sum of each segment; it is initiated with all zeros since, in the beginning, there are no segments due to all elements being 'removed'.

We iterate through removeQueries in reverse order because we want to keep track of the state of the array after each remove operation:

Each time we process an index i from removeQueries, we are effectively re-adding or 'reactivating' an element back into nums and thereby potentially creating a segment or merging it with an adjacent one.
We use the find function to traverse up the parent chain for a given element to find the root parent, applying path compression along the way. Path compression is an optimization that flattens the Union-Find tree, resulting in an almost flat structure over time, which makes subsequent find operations O(1) on average.
The merge function takes two elements a and b, finds their root parents and then unifies the two segments by making one the parent of the other. It then updates the sum of the segment to now include the sum of both merged segments.
After reactivating an element, if it has neighbors that form a positive segment (since only segments of positive integers are considered), we merge this element's segment with the neighbors' segments using merge.
With every reactivation of an element, we could potentially increase the current maximum segment sum if the sum of the newly formed segment is larger. We calculate this using max(mx, s[find(i)]) as mx holds the maximum segment sum found so far.
Finally, before moving to the next removal, we update the ans array with the current maximum segment sum mx. Since we started from the last element, we update ans[j - 1] with mx for each iteration.

By processing the removals in reverse and using Union-Find with path compression, we ensure that the solution performs efficiently for each removal query with each operation taking effectively constant time, leading to an overall time complexity that's almost linear with respect to n.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's consider a small example to illustrate the solution approach. Suppose we have the following input:

nums = [4,2,7,6,9]
removeQueries = [3, 1, 2, 4, 0]

Following our intuition, we want to calculate the maximum segment sum after the i-th element was removed as specified in removeQueries.

Initialize arrays p and s to keep track of the parent indices and the sums of each segment. They will look like this after initialization, where p = [0, 1, 2, 3, 4] and s = [0, 0, 0, 0, 0]. Also, initialize mx = 0 and an answer array ans with the same length as nums, which will store the maximum segment sum after each removal.
Process the removeQueries in reverse:
- First, we 'reactivate' index 0 (nums[0] = 4). No need to merge as it's the first element 'reactivated'. Update s[0] to 4 (the value of nums[0]). The new max segment sum mx becomes 4. Update ans[4] with mx.
Continue reactivating:
- Reactivate index 4 (nums[4] = 9). No adjacent positive segments, so no merges. Update s[4] to 9. Since mx is less than 9, now mx becomes 9. Update ans[3] with mx.
- Reactivate index 2 (nums[2] = 7). No adjacent positive segments, so no merges. Update s[2] to 7. Now mx remains 9 (greater than 7). Update ans[2] with mx.
- Reactivate index 1 (nums[1] = 2). Still no adjacent positive segments, so no merges. Update s[1] to 2. mx remains 9. Update ans[1] with mx.
- Finally, reactivate index 3 (nums[3] = 6). Here we have an adjacent positive segment at index 2 (nums[2] = 7). We call merge on the segments with indices 3 and 2. These segments now form a new segment with a sum of 13. Update s[2] to 13 (as its the root parent after merge). Now mx becomes 13 (greater than the previous max). Update ans[0] with mx.
After processing all removals in reverse, we get our ans array as follows:

[13, 9, 9, 9, 4]

This resulting array represents the maximum segment sum after each i-th removal:

After the 1st removal (index 3), the maximum segment sum is 13.
After the 2nd removal (index 1), the maximum segment sum remains 9.
After the 3rd removal (index 2), the maximum segment sum is still 9.
After the 4th removal (index 4), the maximum segment sum is 9 again.
Finally, after the 5th removal (index 0), the maximum segment sum is 4.

By processing the removals in reverse and merging segments when necessary using the Union-Find data structure with path compression, we efficiently keep track of the maximum segment sums after each removal.

Solution Implementation

1class Solution:
2    def maximumSegmentSum(self, nums: List[int], removeQueries: List[int]) -> List[int]:
3        # Finds the root of the set that element x is part of.
4        def find_root(x):
5            if parent[x] != x:
6                parent[x] = find_root(parent[x])
7            return parent[x]
8
9        # Merges two sets together, connecting components a and b.
10        def merge(a, b):
11            # Find roots of the sets that a and b are part of.
12            root_a, root_b = find_root(a), find_root(b)
13            # Merge the two sets by updating the parent of one root to the other.
14            parent[root_a] = root_b
15            # Update the sum of the merged set to include the sum of both sets.
16            segment_sum[root_b] += segment_sum[root_a]
17
18        # Initialize variables.
19        n = len(nums)
20        parent = list(range(n))   # Parent of each element initially pointing to itself.
21        segment_sum = [0] * n     # Segment sum initialized to 0 for each segment.
22        ans = [0] * n             # Answer list initialized to all zeros.
23        max_segment_sum = 0       # Initialize maximum segment sum to 0.
24
25        # Iterate through the removeQueries in reverse, except the first element.
26        for j in range(n - 1, 0, -1):
27            i = removeQueries[j]
28            segment_sum[i] = nums[i]  # Update the segment sum where the number is added back.
29          
30            # Attempt to merge with the left neighbor if it's part of a segment.
31            if i > 0 and segment_sum[find_root(i - 1)]:
32                merge(i, i - 1)
33            # Attempt to merge with the right neighbor if it's part of a segment.
34            if i < n - 1 and segment_sum[find_root(i + 1)]:
35                merge(i, i + 1)
36            # Update the max_segment_sum with the maximum of itself and the sum of the
37            # segment that contains the current node.
38            max_segment_sum = max(max_segment_sum, segment_sum[find_root(i)])
39            ans[j - 1] = max_segment_sum  # Store the result in the answer list.
40
41        return ans
42

1class Solution {
2    private int[] parent;
3    private long[] segmentSum;
4
5    // Returns the maximum segment sum after each removal
6    public long[] maximumSegmentSum(int[] nums, int[] removeQueries) {
7        int n = nums.length;
8        parent = new int[n];
9        segmentSum = new long[n];
10      
11        // Initialize parent of each element to itself
12        for (int i = 0; i < n; ++i) {
13            parent[i] = i;
14        }
15
16        long[] result = new long[n];
17        long maxSum = 0; // Tracks the current maximum segment sum
18
19        // Process removeQueries in reverse order
20        for (int j = n - 1; j > 0; --j) {
21            int indexToRemove = removeQueries[j];
22            segmentSum[indexToRemove] = nums[indexToRemove]; // Include the value at the removed index
23
24            // Merge with left neighbor if it has a positive sum
25            if (indexToRemove > 0 && segmentSum[find(indexToRemove - 1)] > 0) {
26                merge(indexToRemove, indexToRemove - 1);
27            }
28
29            // Merge with right neighbor if it has a positive sum
30            if (indexToRemove < n - 1 && segmentSum[find(indexToRemove + 1)] > 0) {
31                merge(indexToRemove, indexToRemove + 1);
32            }
33
34            // Update the current maximum segment sum
35            maxSum = Math.max(maxSum, segmentSum[find(indexToRemove)]);
36            result[j - 1] = maxSum;
37        }
38
39        return result;
40    }
41
42    // Finds the root of the set that contains x
43    private int find(int x) {
44        if (parent[x] != x) {
45            parent[x] = find(parent[x]); // Path compression
46        }
47        return parent[x];
48    }
49
50    // Merges sets containing a and b
51    private void merge(int a, int b) {
52        int parentA = find(a);
53        int parentB = find(b);
54        // Make parentB the parent of parentA's set
55        parent[parentA] = parentB;
56        // Sum the values of the two segments
57        segmentSum[parentB] += segmentSum[parentA];
58    }
59}
60

1#include <vector> // Include necessary header for vector
2using namespace std; // Use the standard namespace
3using ll = long long; // Alias for long long type
4
5class Solution {
6public:
7    vector<int> parent; // Represents the parent of each element for disjoint set
8    vector<ll> segmentSum; // Contains the sum of each segment
9
10    // Function to calculate the maximum segment sum after each removal
11    vector<long long> maximumSegmentSum(vector<int>& nums, vector<int>& removeQueries) {
12        int n = nums.size(); // Get the size of input array
13        parent.resize(n); // Resize parent vector to the size of the array
14        for (int i = 0; i < n; ++i) {
15            parent[i] = i; // Initially, set each element's parent to itself
16        }
17        segmentSum.assign(n, 0); // Initialize segment sums to 0
18        vector<ll> answers(n); // Vector to hold the answers
19        ll maxSum = 0; // Variable to hold the maximum sum of any segment
20
21        // Iterate through the remove queries in reverse except for the first element
22        for (int j = n - 1; j > 0; --j) {
23            int indexToRemove = removeQueries[j]; // Index to be removed in this iteration
24            segmentSum[indexToRemove] = nums[indexToRemove]; // Set its segment sum to its value
25
26            // Merge with the previous segment if it exists and has a nonzero sum
27            if (indexToRemove > 0 && segmentSum[find(indexToRemove - 1)] != 0) {
28                merge(indexToRemove, indexToRemove - 1);
29            }
30
31            // Merge with the next segment if it exists and has a nonzero sum
32            if (indexToRemove < n - 1 && segmentSum[find(indexToRemove + 1)] != 0) {
33                merge(indexToRemove, indexToRemove + 1);
34            }
35
36            maxSum = max(maxSum, segmentSum[find(indexToRemove)]); // Update the max sum
37            answers[j - 1] = maxSum; // Set the answer for this query
38        }
39
40        return answers;
41    }
42
43    // Find function for the disjoint-set/union-find structure
44    int find(int x) {
45        if (parent[x] != x) { // Path compression, if x is not the parent of itself
46            parent[x] = find(parent[x]); // Find the actual parent of x
47        }
48        return parent[x]; // Return the parent of x
49    }
50
51    // Merge function for the disjoint-set/union-find structure
52    void merge(int a, int b) {
53        int parentA = find(a); // Find the parent of a
54        int parentB = find(b); // Find the parent of b
55        parent[parentA] = parentB; // Merge the two sets by updating the parent
56        segmentSum[parentB] += segmentSum[parentA]; // Update the segment sum of the new parent
57    }
58};
59

1// TypeScript header imports and aliasing
2type ll = number; // Alias TypeScript's `number` type for 'long long'
3
4let parent: number[]; // Array representing the parent of each element for disjoint set
5let segmentSum: ll[]; // Array containing the sum of each segment
6
7// Function to calculate the maximum segment sum after each removal
8function maximumSegmentSum(nums: number[], removeQueries: number[]): ll[] {
9    const n: number = nums.length; // Get the length of the input array
10    parent = new Array(n); // Initialize the parent array to the size of nums
11    segmentSum = new Array(n).fill(0); // Initialize segment sums with 0
12    const answers: ll[] = new Array(n); // Array to store the answers
13    let maxSum: ll = 0; // Variable to track the maximum sum of any segment
14
15    // Iterate through the remove queries in reverse except for the first element
16    for (let j: number = n - 1; j > 0; --j) {
17        const indexToRemove: number = removeQueries[j]; // Index to be removed in this iteration
18        segmentSum[indexToRemove] = nums[indexToRemove]; // Set the segment sum at the index to its value in nums
19
20        // Merge with the previous segment if it exists and has a nonzero sum
21        if (indexToRemove > 0 && segmentSum[find(indexToRemove - 1)] !== 0) {
22            merge(indexToRemove, indexToRemove - 1);
23        }
24
25        // Merge with the next segment if it exists and has a nonzero sum
26        if (indexToRemove < n - 1 && segmentSum[find(indexToRemove + 1)] !== 0) {
27            merge(indexToRemove, indexToRemove + 1);
28        }
29
30        maxSum = Math.max(maxSum, segmentSum[find(indexToRemove)]); // Update the max sum
31        answers[j - 1] = maxSum; // Set the answer for this query
32    }
33
34    return answers;
35}
36
37// Find function for the disjoint-set/union-find structure
38function find(x: number): number {
39    if (parent[x] !== x) { // Path compression: if x is not the parent of itself
40        parent[x] = find(parent[x]); // Recursively find the actual parent of x
41    }
42    return parent[x]; // Return the parent of x
43}
44
45// Merge function for the disjoint-set/union-find structure
46function merge(a: number, b: number): void {
47    const parentA: number = find(a); // Find the parent of element a
48    const parentB: number = find(b); // Find the parent of element b
49    parent[parentA] = parentB; // Merge the two sets by pointing parent of a to parent of b
50    segmentSum[parentB] += segmentSum[parentA]; // Update the segment sum by adding sum of merged segment
51}
52
53// Initialization as part of global state
54parent = []; // This would be set when `maximumSegmentSum` is called
55segmentSum = []; // Similarly, this is set when `maximumSegmentSum` is called
56

Time and Space Complexity

The provided Python code defines a class Solution with a method maximumSegmentSum that finds the maximum segment sum after each removal query.

Time Complexity:

The main operations of the algorithm are performed inside a loop that runs n - 1 times, where n is the length of the nums array. This loop handles the updating of segment sums and parent pointers for the union-find data structure.

For each iteration:

The find function involves path compression and may take O(log n) time in the worst case if no path compression was used, but with path compression, the amortized time complexity is O(alpha(n)), where alpha(n) is the inverse Ackermann function which grows very slowly and for all practical purposes is less than 5.
The merge function invokes the find function twice and does constant-time work for the merge, which means it also operates in O(alpha(n)) time.
Updating the maximum segment sum using max is an O(1) operation since it's a simple comparison and assignment.

Given the above, the total time complexity is O(n * alpha(n)), where n is the number of elements in nums.

Space Complexity:

The auxiliary space used by the algorithm includes:

The parent array p of length n, contributing O(n) space.
The segment sum array s also of length n, contributing another O(n) space.
The answer array ans, again of length n, contributing O(n) space.

Therefore, the total space complexity of the algorithm is O(n).

In conclusion, the time complexity is O(n * alpha(n)), and the space complexity is O(n).

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Question 1 out of 10

Depth first search is equivalent to which of the tree traversal order?

Pre-order traversal

Post-order traversal

In-order traversal

Level-order traversal

I don’t know