Skip to content

2382. Maximum Segment Sum After Removals

Difficulty Topics

Description

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.

 

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 <= 105
  • 1 <= nums[i] <= 109
  • 0 <= removeQueries[i] < n
  • All the values of removeQueries are unique.

Solution

maximum-segment-sum-after-removals.py
class UnionFind:
    def __init__(self):
        self._parent = {}
        self.segmentSum = {}
        self.maxSegmentSum = 0

    def merge(self, u, value):
        self._parent[u] = u
        self.segmentSum[u] = value
        self.maxSegmentSum = max(self.maxSegmentSum, value)

        if u - 1 in self._parent:
            self.union(u, u - 1)

        if u + 1 in self._parent:
            self.union(u, u + 1)

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a == b:
            return

        self.segmentSum[a] += self.segmentSum[b]
        self._parent[b] = a
        self.maxSegmentSum = max(self.maxSegmentSum, self.segmentSum[a])

    def find(self, x):
        if x != self._parent[x]:
            self._parent[x] = self.find(self._parent[x])

        return self._parent[x]

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

        for i in range(n - 1, -1, -1):
            res[i] = uf.maxSegmentSum
            u = removeQueries[i]
            uf.merge(u, nums[u])

        return res