Maximum Segment Sum After Removals

Problem

You are given two 0-indexed integer arrays nums and removeQueries, both of length n. For the ith query, the element in nums at the index removeQueries[i] is removed, splitting nums into different segments.

A segment is a contiguous sequence of positive integers in nums. A segment sum is the sum of every element in a segment.

Return an integer arrayanswer , of lengthn , whereanswer[i]_is themaximum segment sum after applying the _ith removal.

Note: The same index will not be removed more than once.

Examples

Example 1

    
    
    Input: nums = [1,2,5,6,1], removeQueries = [0,3,2,4,1]
    Output: [14,7,2,2,0]
    Explanation: Using 0 to indicate a removed element, the answer is as follows:
    Query 1: Remove the 0th element, nums becomes [0,2,5,6,1] and the maximum segment sum is 14 for segment [2,5,6,1].
    Query 2: Remove the 3rd element, nums becomes [0,2,5,0,1] and the maximum segment sum is 7 for segment [2,5].
    Query 3: Remove the 2nd element, nums becomes [0,2,0,0,1] and the maximum segment sum is 2 for segment [2]. 
    Query 4: Remove the 4th element, nums becomes [0,2,0,0,0] and the maximum segment sum is 2 for segment [2]. 
    Query 5: Remove the 1st element, nums becomes [0,0,0,0,0] and the maximum segment sum is 0, since there are no segments.
    Finally, we return [14,7,2,2,0].

Example 2

    
    
    Input: nums = [3,2,11,1], removeQueries = [3,2,1,0]
    Output: [16,5,3,0]
    Explanation: Using 0 to indicate a removed element, the answer is as follows:
    Query 1: Remove the 3rd element, nums becomes [3,2,11,0] and the maximum segment sum is 16 for segment [3,2,11].
    Query 2: Remove the 2nd element, nums becomes [3,2,0,0] and the maximum segment sum is 5 for segment [3,2].
    Query 3: Remove the 1st element, nums becomes [3,0,0,0] and the maximum segment sum is 3 for segment [3].
    Query 4: Remove the 0th element, nums becomes [0,0,0,0] and the maximum segment sum is 0, since there are no segments.
    Finally, we return [16,5,3,0].
    

Constraints

• n == nums.length == removeQueries.length
• 1 <= n <= 10^5
• 1 <= nums[i] <= 10^9
• 0 <= removeQueries[i] < n
• All the values of removeQueries are unique.

Solution

Method 1 – Union-Find with Segment Tracking

Intuition

We process removals in reverse (i.e., we "add back" elements), using union-find to dynamically merge segments and track their sums. This allows us to efficiently maintain the maximum segment sum after each removal.

Approach

1. Initialize all elements as removed (inactive).
2. Process removeQueries in reverse, each time activating the removed index.
3. Use union-find to merge adjacent active segments and maintain their sums.
4. After each addition, update the current maximum segment sum.
5. Store the result for each step, then reverse the answer at the end.

Code

C++

class Solution {
public:
    vector<long long> maximumSegmentSum(vector<int>& nums, vector<int>& removeQueries) {
        int n = nums.size();
        vector<long long> ans(n), segSum(n), parent(n);
        vector<bool> active(n, false);
        long long mx = 0;
        iota(parent.begin(), parent.end(), 0);
        function<int(int)> find = [&](int x) {
            return parent[x] == x ? x : parent[x] = find(parent[x]);
        };
        for (int i = n - 1; i > 0; --i) {
            int idx = removeQueries[i];
            active[idx] = true;
            segSum[idx] = nums[idx];
            if (idx > 0 && active[idx - 1]) {
                int l = find(idx - 1);
                parent[idx] = l;
                segSum[l] += segSum[idx];
                segSum[idx] = 0;
                idx = l;
            }
            if (idx + 1 < n && active[idx + 1]) {
                int r = find(idx + 1);
                parent[r] = idx;
                segSum[idx] += segSum[r];
                segSum[r] = 0;
            }
            mx = max(mx, segSum[find(idx)]);
            ans[i - 1] = mx;
        }
        return ans;
    }
};

Go

func maximumSegmentSum(nums []int, removeQueries []int) []int64 {
    n := len(nums)
    ans := make([]int64, n)
    segSum := make([]int64, n)
    parent := make([]int, n)
    active := make([]bool, n)
    for i := range parent {
        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]
    }
    var mx int64
    for i := n - 1; i > 0; i-- {
        idx := removeQueries[i]
        active[idx] = true
        segSum[idx] = int64(nums[idx])
        if idx > 0 && active[idx-1] {
            l := find(idx - 1)
            parent[idx] = l
            segSum[l] += segSum[idx]
            segSum[idx] = 0
            idx = l
        }
        if idx+1 < n && active[idx+1] {
            r := find(idx + 1)
            parent[r] = idx
            segSum[idx] += segSum[r]
            segSum[r] = 0
        }
        if segSum[find(idx)] > mx {
            mx = segSum[find(idx)]
        }
        ans[i-1] = mx
    }
    return ans
}

Java

class Solution {
    public long[] maximumSegmentSum(int[] nums, int[] removeQueries) {
        int n = nums.length;
        long[] ans = new long[n], segSum = new long[n];
        int[] parent = new int[n];
        boolean[] active = new boolean[n];
        for (int i = 0; i < n; ++i) parent[i] = i;
        long mx = 0;
        for (int i = n - 1; i > 0; --i) {
            int idx = removeQueries[i];
            active[idx] = true;
            segSum[idx] = nums[idx];
            if (idx > 0 && active[idx - 1]) {
                int l = find(parent, idx - 1);
                parent[idx] = l;
                segSum[l] += segSum[idx];
                segSum[idx] = 0;
                idx = l;
            }
            if (idx + 1 < n && active[idx + 1]) {
                int r = find(parent, idx + 1);
                parent[r] = idx;
                segSum[idx] += segSum[r];
                segSum[r] = 0;
            }
            mx = Math.max(mx, segSum[find(parent, idx)]);
            ans[i - 1] = mx;
        }
        return ans;
    }
    private int find(int[] parent, int x) {
        if (parent[x] != x) parent[x] = find(parent, parent[x]);
        return parent[x];
    }
}

Kotlin

class Solution {
    fun maximumSegmentSum(nums: IntArray, removeQueries: IntArray): LongArray {
        val n = nums.size
        val ans = LongArray(n)
        val segSum = LongArray(n)
        val parent = IntArray(n) { it }
        val active = BooleanArray(n)
        var mx = 0L
        fun find(x: Int): Int = if (parent[x] == x) x else find(parent[x]).also { parent[x] = it }
        for (i in n - 1 downTo 1) {
            var idx = removeQueries[i]
            active[idx] = true
            segSum[idx] = nums[idx].toLong()
            if (idx > 0 && active[idx - 1]) {
                val l = find(idx - 1)
                parent[idx] = l
                segSum[l] += segSum[idx]
                segSum[idx] = 0
                idx = l
            }
            if (idx + 1 < n && active[idx + 1]) {
                val r = find(idx + 1)
                parent[r] = idx
                segSum[idx] += segSum[r]
                segSum[r] = 0
            }
            mx = maxOf(mx, segSum[find(idx)])
            ans[i - 1] = mx
        }
        return ans
    }
}

Python

class Solution:
    def maximumSegmentSum(self, nums: list[int], removeQueries: list[int]) -> list[int]:
        n = len(nums)
        ans: list[int] = [0] * n
        seg_sum: list[int] = [0] * n
        parent: list[int] = list(range(n))
        active: list[bool] = [False] * n
        mx = 0
        def find(x: int) -> int:
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]
        for i in range(n - 1, 0, -1):
            idx = removeQueries[i]
            active[idx] = True
            seg_sum[idx] = nums[idx]
            if idx > 0 and active[idx - 1]:
                l = find(idx - 1)
                parent[idx] = l
                seg_sum[l] += seg_sum[idx]
                seg_sum[idx] = 0
                idx = l
            if idx + 1 < n and active[idx + 1]:
                r = find(idx + 1)
                parent[r] = idx
                seg_sum[idx] += seg_sum[r]
                seg_sum[r] = 0
            mx = max(mx, seg_sum[find(idx)])
            ans[i - 1] = mx
        return ans

Rust

impl Solution {
    pub fn maximum_segment_sum(nums: Vec<i32>, remove_queries: Vec<i32>) -> Vec<i64> {
        let n = nums.len();
        let mut ans = vec![0i64; n];
        let mut seg_sum = vec![0i64; n];
        let mut parent: Vec<usize> = (0..n).collect();
        let mut active = vec![false; n];
        let mut mx = 0i64;
        fn find(parent: &mut Vec<usize>, x: usize) -> usize {
            if parent[x] != x {
                parent[x] = find(parent, parent[x]);
            }
            parent[x]
        }
        for i in (1..n).rev() {
            let mut idx = remove_queries[i] as usize;
            active[idx] = true;
            seg_sum[idx] = nums[idx] as i64;
            if idx > 0 && active[idx - 1] {
                let l = find(&mut parent, idx - 1);
                parent[idx] = l;
                seg_sum[l] += seg_sum[idx];
                seg_sum[idx] = 0;
                idx = l;
            }
            if idx + 1 < n && active[idx + 1] {
                let r = find(&mut parent, idx + 1);
                parent[r] = idx;
                seg_sum[idx] += seg_sum[r];
                seg_sum[r] = 0;
            }
            mx = mx.max(seg_sum[find(&mut parent, idx)]);
            ans[i - 1] = mx;
        }
        ans
    }
}

TypeScript

class Solution {
    maximumSegmentSum(nums: number[], removeQueries: number[]): number[] {
        const n = nums.length;
        const ans: number[] = Array(n).fill(0);
        const segSum: number[] = Array(n).fill(0);
        const parent: number[] = Array.from({length: n}, (_, i) => i);
        const active: boolean[] = Array(n).fill(false);
        let mx = 0;
        const find = (x: number): number => parent[x] === x ? x : (parent[x] = find(parent[x]));
        for (let i = n - 1; i > 0; --i) {
            let idx = removeQueries[i];
            active[idx] = true;
            segSum[idx] = nums[idx];
            if (idx > 0 && active[idx - 1]) {
                const l = find(idx - 1);
                parent[idx] = l;
                segSum[l] += segSum[idx];
                segSum[idx] = 0;
                idx = l;
            }
            if (idx + 1 < n && active[idx + 1]) {
                const r = find(idx + 1);
                parent[r] = idx;
                segSum[idx] += segSum[r];
                segSum[r] = 0;
            }
            mx = Math.max(mx, segSum[find(idx)]);
            ans[i - 1] = mx;
        }
        return ans;
    }
}

Complexity

• ⏰ Time complexity: O(n α(n)), where α(n) is the inverse Ackermann function for union-find. Each union/find operation is nearly constant time, and we do O(n) operations.
• 🧺 Space complexity: O(n), for the parent, segment sum, and active arrays.