LeetCode 2382 - Maximum Segment Sum After Removals

LeetCode Problem 2382

Difficulty: 🔴 Hard
Topics: Array, Union-Find, Prefix Sum, Ordered Set

Solution

Problem Understanding

The problem gives us two arrays, nums and removeQueries, both of length n.

The nums array contains positive integers. Initially, all elements are present, forming one contiguous segment. Then, elements are removed one by one according to the order defined in removeQueries.

After each removal, the array becomes split into one or more contiguous segments of remaining positive integers. Removed elements effectively act as separators. For every removal operation, we must determine the maximum segment sum among all currently valid contiguous segments.

The goal is to return an array answer, where answer[i] represents the largest segment sum after performing the i-th removal.

To understand the process more clearly, suppose we remove an element in the middle of an existing segment. That segment breaks into two smaller segments. If we remove an element at an edge, the segment simply shrinks. After each such modification, we must efficiently determine which remaining segment has the largest total sum.

The constraints are the most important clue for choosing an algorithm:

• n can be as large as 10^5
• nums[i] can be as large as 10^9
• All removal indices are unique

A straightforward simulation after every removal would be too expensive because repeatedly recomputing segment sums can lead to quadratic complexity.

An important observation is that all numbers are positive. This means larger segments always have larger or equal sums, and segment sums can be maintained incrementally without worrying about negative values reducing totals.

There are several edge cases to keep in mind:

• When all elements are removed, the answer must become 0 because no segments remain.
• Removing from the beginning or end of a segment behaves differently than removing from the middle.
• A single remaining element forms a valid segment by itself.
• Since nums[i] can be up to 10^9, segment sums may exceed 32-bit integer limits, so we must use 64-bit integers.

Approaches

Brute Force Approach

A naive solution would directly simulate the removal process.

After each removal:

1. Mark the removed index.
2. Scan the entire array from left to right.
3. Identify contiguous segments of non-removed elements.
4. Compute the sum of every segment.
5. Track the maximum segment sum.

This approach works because it exactly follows the problem definition. Every query recomputes the current state of the array and evaluates all segments.

However, it is too slow.

For each removal, we may scan the entire array of length n. Since there are n removals, the total complexity becomes O(n²).

With n = 10^5, this could require around 10^10 operations, which is far beyond practical limits.

Key Insight for the Optimal Approach

The difficult part is that removing elements splits segments, which is hard to maintain efficiently.

Instead of thinking forward, we can think backward.

Rather than removing elements one at a time, we can start with an empty array and add elements back in reverse order.

This transforms a difficult splitting problem into an easier merging problem.

Why is merging easier?

When we add an element back:

• It starts as a new segment.
• If its left neighbor exists, we merge with the left segment.
• If its right neighbor exists, we merge with the right segment.

Efficient merging is exactly what Union-Find (Disjoint Set Union, DSU) is designed for.

We maintain:

• A parent array for DSU structure
• A segment_sum array storing each component's total sum
• An active array to track whether an index exists
• A running max_segment_sum

At every reversed step:

1. Add the element back.
2. Merge adjacent active components.
3. Update the maximum segment sum.

Because each union and find operation is nearly constant time, the solution becomes efficient enough for 10^5 elements.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n²)	O(n)	Recompute all segments after every removal
Optimal	O(n α(n))	O(n)	Reverse process with Union-Find merging

Here, α(n) is the inverse Ackermann function, which grows so slowly that it behaves like a constant in practice.

Algorithm Walkthrough

Optimal Algorithm Using Reverse Processing + Union-Find

1. Initialize the answer array with size n.

We will fill it in reverse order because we are processing additions instead of removals. 2. Create a Union-Find structure.

We maintain:

• parent[i], the representative of a component
• segment_sum[i], the sum of the connected component rooted at i
• active[i], whether index i currently exists

Initially, all elements are inactive because we start from a fully removed array. 3. Maintain a variable max_sum.

This stores the maximum segment sum currently present.

Initially:

max_sum = 0

because no segments exist. 4. Process removeQueries in reverse order.

Before restoring an element, record the current max_sum into the answer array.

Why?

Because reversing removals means the current state corresponds exactly to the array after removal in the original order. 5. Activate the current index.

Let:

idx = removeQueries[i]

Mark it active and initialize:

segment_sum[idx] = nums[idx]
parent[idx] = idx

This forms a single-element segment. 6. Check the left neighbor.

If idx - 1 exists and is active, merge the two components.

This restores continuity between adjacent elements. 7. Check the right neighbor.

If idx + 1 exists and is active, merge again.

This ensures larger segments are rebuilt incrementally. 8. Update the maximum segment sum.

After merging, retrieve the root and update:

max_sum = max(max_sum, component_sum)

9. Continue until all elements are restored.

Since we processed in reverse, the answer array already contains results in the correct order.

Why it works

The key invariant is that after processing the reverse step i, the Union-Find structure exactly represents the array state after the i-th removal in the original order.

Processing removals directly causes segments to split, which is difficult to maintain efficiently. By reversing the process, every operation becomes a merge between neighboring segments, which Union-Find handles naturally.

Because every merge preserves the exact segment sums and max_sum always stores the largest connected component sum, the algorithm always produces the correct answer.

Python Solution

from typing import List

class Solution:
    def maximumSegmentSum(
        self,
        nums: List[int],
        removeQueries: List[int]
    ) -> List[int]:
        n = len(nums)

        parent = list(range(n))
        segment_sum = [0] * n
        active = [False] * n
        answer = [0] * n

        def find(x: int) -> int:
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        def union(x: int, y: int) -> int:
            root_x = find(x)
            root_y = find(y)

            if root_x == root_y:
                return root_x

            parent[root_y] = root_x
            segment_sum[root_x] += segment_sum[root_y]

            return root_x

        max_sum = 0

        for i in range(n - 1, -1, -1):
            answer[i] = max_sum

            index = removeQueries[i]

            active[index] = True
            segment_sum[index] = nums[index]

            if index > 0 and active[index - 1]:
                union(index, index - 1)

            if index < n - 1 and active[index + 1]:
                union(index, index + 1)

            root = find(index)
            max_sum = max(max_sum, segment_sum[root])

        return answer

The implementation closely follows the algorithm.

We first initialize the Union-Find data structures. The parent array stores component representatives, while segment_sum tracks the total sum for each connected segment.

The find function uses path compression to speed up future lookups by flattening the tree structure.

The union function merges two neighboring components and updates the combined segment sum. Since we only care about total sums, we simply attach one root to the other.

The main loop processes removeQueries in reverse. Before restoring an element, we store the current maximum segment sum into answer[i]. This corresponds to the array state after removal in the original sequence.

Each restored index starts as a one-element segment. We then merge with active neighbors if they exist, rebuild the correct segment structure, and update the running maximum.

Finally, we return the completed answer array.

Go Solution

func maximumSegmentSum(nums []int, removeQueries []int) []int64 {
	n := len(nums)

	parent := make([]int, n)
	segmentSum := make([]int64, n)
	active := make([]bool, n)
	answer := make([]int64, n)

	for i := 0; i < n; i++ {
		parent[i] = i
	}

	var find func(int) int
	find = func(x int) int {
		if parent[x] != x {
			parent[x] = find(parent[x])
		}
		return parent[x]
	}

	union := func(x, y int) int {
		rootX := find(x)
		rootY := find(y)

		if rootX == rootY {
			return rootX
		}

		parent[rootY] = rootX
		segmentSum[rootX] += segmentSum[rootY]

		return rootX
	}

	var maxSum int64 = 0

	for i := n - 1; i >= 0; i-- {
		answer[i] = maxSum

		index := removeQueries[i]

		active[index] = true
		segmentSum[index] = int64(nums[index])

		if index > 0 && active[index-1] {
			union(index, index-1)
		}

		if index < n-1 && active[index+1] {
			union(index, index+1)
		}

		root := find(index)
		if segmentSum[root] > maxSum {
			maxSum = segmentSum[root]
		}
	}

	return answer
}

The Go implementation is very similar to the Python version, but there are some language-specific differences.

The biggest difference is integer handling. Since segment sums may exceed 32-bit limits, we use int64 for segmentSum, maxSum, and the return type.

Go also uses slices instead of Python lists, and recursion is used for path compression in find.

Unlike Python, Go does not have built-in arbitrary precision integers, so using int64 is necessary to avoid overflow.

Worked Examples

Example 1

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

We process in reverse:

Reverse order: [1,4,2,3,0]

Step	Added Index	Active Segments	Segment Sums	Max Sum	Answer Written
Start	-	[]	[]	0	-
i=4	1	[2]	[2]	2	answer[4]=0
i=3	4	[2], [1]	[2,1]	2	answer[3]=2
i=2	2	[2,5] and [1]	[7,1]	7	answer[2]=2
i=1	3	[2,5,6,1]	[14]	14	answer[1]=7
i=0	0	[1,2,5,6,1]	[15]	15	answer[0]=14

Final answer:

[14,7,2,2,0]

Example 2

nums = [3,2,11,1]
removeQueries = [3,2,1,0]

Reverse order:

[0,1,2,3]

Step	Added Index	Active Segments	Segment Sums	Max Sum	Answer Written
Start	-	[]	[]	0	-
i=3	0	[3]	[3]	3	answer[3]=0
i=2	1	[3,2]	[5]	5	answer[2]=3
i=1	2	[3,2,11]	[16]	16	answer[1]=5
i=0	3	[3,2,11,1]	[17]	17	answer[0]=16

Final answer:

[16,5,3,0]

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n α(n))	Each index is activated once, DSU operations are nearly constant
Space	O(n)	Arrays for parent, active state, segment sums, and answer

The runtime comes from processing every element exactly once and performing at most two union operations per element. Thanks to path compression in Union-Find, find and union become nearly constant time, giving an effective complexity of O(n α(n)).

The space complexity is linear because we store several arrays proportional to n.

Test Cases

class Solution:
    def maximumSegmentSum(self, nums, removeQueries):
        n = len(nums)

        parent = list(range(n))
        segment_sum = [0] * n
        active = [False] * n
        answer = [0] * n

        def find(x):
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        def union(x, y):
            root_x = find(x)
            root_y = find(y)

            if root_x != root_y:
                parent[root_y] = root_x
                segment_sum[root_x] += segment_sum[root_y]

        max_sum = 0

        for i in range(n - 1, -1, -1):
            answer[i] = max_sum

            idx = removeQueries[i]
            active[idx] = True
            segment_sum[idx] = nums[idx]

            if idx > 0 and active[idx - 1]:
                union(idx, idx - 1)

            if idx < n - 1 and active[idx + 1]:
                union(idx, idx + 1)

            root = find(idx)
            max_sum = max(max_sum, segment_sum[root])

        return answer

sol = Solution()

assert sol.maximumSegmentSum(
    [1,2,5,6,1],
    [0,3,2,4,1]
) == [14,7,2,2,0]  # Example 1

assert sol.maximumSegmentSum(
    [3,2,11,1],
    [3,2,1,0]
) == [16,5,3,0]  # Example 2

assert sol.maximumSegmentSum(
    [5],
    [0]
) == [0]  # Single element

assert sol.maximumSegmentSum(
    [1,1,1,1],
    [0,1,2,3]
) == [3,2,1,0]  # Sequential removals

assert sol.maximumSegmentSum(
    [10,20,30],
    [1,0,2]
) == [30,30,0]  # Middle split first

assert sol.maximumSegmentSum(
    [1000000000,1000000000],
    [0,1]
) == [1000000000,0]  # Large values

assert sol.maximumSegmentSum(
    [4,5,6,7],
    [3,2,1,0]
) == [15,9,4,0]  # Reverse removals

assert sol.maximumSegmentSum(
    [8,1,2,9],
    [1,2,0,3]
) == [17,9,9,0]  # Multiple merges and splits

Test	Why
Example 1	Validates the primary example
Example 2	Verifies reverse reconstruction
Single element	Smallest valid input
Sequential removals	Tests shrinking from left
Middle split first	Tests segment splitting behavior
Large values	Verifies 64-bit safety
Reverse removals	Tests repeated segment growth
Multiple merges and splits	Ensures neighbor merging works correctly

Edge Cases

Single Element Array

When n = 1, removing the only element leaves no segments.

This case is easy to mishandle because some implementations assume neighboring elements exist. Our implementation safely checks bounds before attempting unions, so no invalid access occurs.

Example:

nums = [5]
removeQueries = [0]

Output:

[0]

Removing Elements That Split Large Segments

Removing an element in the middle of a segment creates two disconnected parts.

A forward simulation must carefully maintain split segments, which is error-prone. Our reverse-processing approach avoids splitting entirely. Instead, segments are only merged, which Union-Find handles efficiently.

Example:

nums = [10,20,30]
removeQueries = [1,0,2]

Removing index 1 splits [10,20,30] into [10] and [30].

Very Large Numbers

Since nums[i] can be 10^9 and there may be 10^5 elements, total segment sums can reach:

10^14

This exceeds 32-bit integer limits.

The implementation handles this correctly by using Python's arbitrary precision integers and int64 in Go. This prevents overflow and ensures correct results for large inputs.