Maximum segment sum after removals

Array Union-Find Prefix Sum Ordered Set

49.3% acceptance

Feb 25, 2026

497

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 array answer, of length n, where answer[i] is the maximum segment sum after applying the ith removal. Note: The same index will not be removed more than once.

Solution

Rust

Time O(2^n)

Space O(n)

LeetCode

solution.rs

fn find_uf(parent: &mut Vec<usize>, mut x: usize) -> usize {
  while parent[x] != x { parent[x] = parent[parent[x]]; x = parent[x]; }
  x
}

impl Solution {
  pub fn maximum_segment_sum(nums: Vec<i32>, remove_queries: Vec<i32>) -> Vec<i64> {
    let n = nums.len();
    let mut parent: Vec<usize> = (0..n).collect();
    let mut seg_sum = vec![0i64; n];
    let mut max_sum = 0i64;
    let mut present = vec![false; n];
    let mut ans = vec![0i64; n];
    for i in (0..n).rev() {
      let idx = remove_queries[i] as usize;
      ans[i] = max_sum;
      present[idx] = true;
      parent[idx] = idx;
      seg_sum[idx] = nums[idx] as i64;
      if idx + 1 < n && present[idx + 1] {
        let ri = find_uf(&mut parent, idx + 1);
        seg_sum[ri] += seg_sum[idx];
        parent[idx] = ri;
      }
      if idx > 0 && present[idx - 1] {
        let li = find_uf(&mut parent, idx - 1);
        let me = find_uf(&mut parent, idx);
        seg_sum[li] += seg_sum[me];
        parent[me] = li;
      }
      let root = find_uf(&mut parent, idx);
      max_sum = max_sum.max(seg_sum[root]);
    }
    ans
  }
}