2382

2382. Maximum Segment Sum After Removals (Hard)

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.

Companies: Infosys

Related Topics:
Array, Union Find, Prefix Sum, Ordered Set

Solution 1. Manage Intervals with Ordered Set

// OJ: https://leetcode.com/problems/maximum-segment-sum-after-removals
// Author: github.com/lzl124631x
// Time: O(NlogN)
// Space: O(N)
class Solution {
public:
    vector<long long> maximumSegmentSum(vector<int>& A, vector<int>& Q) {
        int N = A.size();
        vector<long long> ans;
        typedef pair<int, int> PII;
        vector<long long> sum(N + 1);
        for (int i = 0; i < N; ++i) sum[i + 1] = sum[i] + A[i];
        auto cmp = [&](auto &a, auto &b) { return a.second < b.second; };
        set<PII, decltype(cmp)> intervals(cmp);
        intervals.insert({0, N - 1});
        multiset<long long> s{sum.back()};
        for (int q : Q) {
            auto it = intervals.lower_bound({q, q});
            auto [a, b] = *it;
            intervals.erase(it);
            s.erase(s.find(sum[b + 1] - sum[a]));
            if (a <= q - 1) {
                intervals.insert({a, q - 1});
                s.insert(sum[q] - sum[a]);
            }
            if (b >= q + 1) {
                intervals.insert({q + 1, b});
                s.insert(sum[b + 1] - sum[q + 1]);
            }
            ans.push_back(s.size() ? *s.rbegin() : 0);
        }
        return ans;
    }
};

Solution 2. Union Find

Traverse the queries in the reverse order. Use Union find to keep track of the segment sums.

// OJ: https://leetcode.com/problems/maximum-segment-sum-after-removals
// Author: github.com/lzl124631x
// Time: O(NlogN)
// Space: O(N)
class Solution {
public:
    vector<long long> maximumSegmentSum(vector<int>& A, vector<int>& Q) {
        long long N = A.size(), mxSum = 0;
        vector<long long> ans(N), sum(begin(A), end(A));
        vector<int> id(N), seen(N);
        iota(begin(id), end(id), 0);
        function<int(int)> find = [&](int a) {
            return id[a] == a ? a : (id[a] = find(id[a]));
        };
        auto connect = [&](int a, int b) {
            int p = find(a), q = find(b);
            if (p == q) return;
            id[q] = p;
            sum[p] += sum[q];
        };
        for (int i = N - 1; i >= 0; --i) {
            int q = Q[i];
            seen[q] = 1;
            if (q - 1 >= 0 && seen[q - 1]) connect(q - 1, q);
            if (q + 1 < N && seen[q + 1]) connect(q, q + 1);
            mxSum = max(mxSum, sum[find(q)]);
            if (i > 0) ans[i - 1] = mxSum;
        }
        return ans;
    }
};