2382. Maximum Segment Sum After Removals

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  vector<long long> maximumSegmentSum(vector<int>& nums,
                                      vector<int>& removeQueries) {
    const int n = nums.size();
    long maxSum = 0;
    vector<long long> ans(n);
    // For the segment [l, r], record its sum in sum[l] and sum[r]
    vector<long long> sum(n);
    // For the segment [l, r], record its count in count[l] and count[r]
    vector<int> count(n);

    for (int i = n - 1; i >= 0; --i) {
      ans[i] = maxSum;
      const int j = removeQueries[i];

      // Calculate `segmentSum`.
      const long leftSum = j > 0 ? sum[j - 1] : 0;
      const long rightSum = j + 1 < n ? sum[j + 1] : 0;
      const long segmentSum = nums[j] + leftSum + rightSum;

      // Calculate `segmentCount`.
      const int leftCount = j > 0 ? count[j - 1] : 0;
      const int rightCount = j + 1 < n ? count[j + 1] : 0;
      const int segmentCount = 1 + leftCount + rightCount;

      // Update the sum and count of the segment [l, r].
      const int l = j - leftCount;
      const int r = j + rightCount;
      sum[l] = segmentSum;
      sum[r] = segmentSum;
      count[l] = segmentCount;
      count[r] = segmentCount;
      maxSum = max(maxSum, segmentSum);
    }

    return ans;
  }
};

class Solution {
  public long[] maximumSegmentSum(int[] nums, int[] removeQueries) {
    final int n = nums.length;
    long maxSum = 0;
    long[] ans = new long[n];
    // For the segment [l, r], record its sum in sum[l] and sum[r]
    long[] sum = new long[n];
    // For the segment [l, r], record its count in count[l] and count[r]
    int[] count = new int[n];

    for (int i = n - 1; i >= 0; --i) {
      ans[i] = maxSum;
      final int j = removeQueries[i];

      // Calculate `segmentSum`.
      final long leftSum = j > 0 ? sum[j - 1] : 0;
      final long rightSum = j + 1 < n ? sum[j + 1] : 0;
      final long segmentSum = nums[j] + leftSum + rightSum;

      // Calculate `segmentCount`.
      final int leftCount = j > 0 ? count[j - 1] : 0;
      final int rightCount = j + 1 < n ? count[j + 1] : 0;
      final int segmentCount = 1 + leftCount + rightCount;

      // Update the sum and count of the segment [l, r].
      final int l = j - leftCount;
      final int r = j + rightCount;
      sum[l] = segmentSum;
      sum[r] = segmentSum;
      count[l] = segmentCount;
      count[r] = segmentCount;
      maxSum = Math.max(maxSum, segmentSum);
    }

    return ans;
  }
}

class Solution:
  def maximumSegmentSum(
      self,
      nums: list[int],
      removeQueries: list[int],
  ) -> list[int]:
    n = len(nums)
    maxSum = 0
    ans = [0] * n
    # For the segment [l, r], record its sum in summ[l] and summ[r]
    summ = [0] * n
    # For the segment [l, r], record its count in count[l] and count[r]
    count = [0] * n

    for i in reversed(range(n)):
      ans[i] = maxSum
      j = removeQueries[i]

      # Calculate `segmentSum`.
      leftSum = summ[j - 1] if j > 0 else 0
      rightSum = summ[j + 1] if j + 1 < n else 0
      segmentSum = nums[j] + leftSum + rightSum

      # Calculate `segmentCount`.
      leftCount = count[j - 1] if j > 0 else 0
      rightCount = count[j + 1] if j + 1 < n else 0
      segmentCount = 1 + leftCount + rightCount

      # Update `summ` and `count` of the segment [l, r].
      l = j - leftCount
      r = j + rightCount
      summ[l] = segmentSum
      summ[r] = segmentSum
      count[l] = segmentCount
      count[r] = segmentCount
      maxSum = max(maxSum, segmentSum)

    return ans