Maximum Segment Sum After Removals

Maximum Segment Sum After Removals

You are given two 0-indexed integer arrays $nums$ and $removeQueries$, both of length $n$. For the $i^{th}$ 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 $i^{th}$ 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 \leq n \leq {10}^{5}$
$1 \leq nums[i] \leq {10}^{9}$
$0 \leq removeQueries[i] < n$
All the values of $removeQueries$ are unique.

 

解题思路

　　这题可以正着做（在线做法），用平衡树（这里用STL的set和multiset实现）。翻着做（离线做法），用并查集。

　　比赛的时候是正着做，思路有了并且正确，但死在不熟悉STL和边界情况上了。

　　正着做就是用一个集合set来维护若干组区间，每一次操作就相当于将某个区间裂开成两个区间（也有可能是一个区间）。先根据当前要删除的下标$x$找到覆盖这个下标的区间$(left, right)$，这里可以通过set::lower_bound()来实现（这里有个边界要处理，当时找了半天bug都没找到，解析在代码下方）。然后把区间裂成$(left, x-1)$和$(x+1, right)$（当然也可能只有一个区间），同时把区间$(left, right)$删除。同时还要开一个multiset来维护set中各个区间的总和，multiset::rbegin()就是所有区间总和的最大值，这里是把multiset当作堆来使用（当时比赛的时候想到用优先队列来维护，但优先队列不支持删除操作，所以还写了些其他东西来补助这个操作，写得很乱）。

　　AC代码如下：

class Solution {
public:
    vector<long long> maximumSegmentSum(vector<int>& nums, vector<int>& removeQueries) {
        int n = nums.size();
        vector<long long> s(n + 1); // 前缀和数组
        for (int i = 1; i <= n; i++) {
            s[i] = s[i - 1] + nums[i - 1];
        }

        set<pair<int, int>> st;     // 维护区间
        st.insert({1, n});
        multiset<long long> mst;    // 维护各个区间的总和，可以得到最大值
        mst.insert(s[n]);
        vector<long long> ans;

        for (int i = 0; i < n; i++) {
            int x = removeQueries[i] + 1;
            auto it = --st.upper_bound({x, n + 1}); // 找到可以覆盖x的区间。{x, n+1}的第二维相当于正无穷，找到左端点严格大于x的区间，那么前一个区间的左端点就一定<=x

            if (x >= it->first) {   // 裂成左边部分的区间(left, x-1)
                st.insert({x + 1, it->second});
                mst.insert(s[it->second] - s[x]);
            }
            if (x <= it->second) {  // 裂成右边部分的区间(x+1, right)
                st.insert({it->first, x - 1});
                mst.insert(s[x - 1] - s[it->first - 1]);
            }

            mst.erase(mst.find(s[it->second] - s[it->first - 1]));  // 删去mst中区间(left, right)的值，这里要删除迭代器，如果直接删除对应的值，就会把所有相同的值都删除而不是只删这一个
            st.erase(it);   // 删除区间(left right)

            ans.push_back(*mst.rbegin());
        }

        return ans;
    }
};

　　这里记录一个bug。就是在上面代码中的$18$行upper_bound()操作，一开始我传入的值是$\{ {x,x} \}$，然后在跑某组数据时发生执行出错。原因就是，如果此时set中只有一个区间$(1,2)$，而我传入的是$\{ {1,1} \}$，也就是找到严格大于$(1,1)$的pair，可以发现第一维$1=1$，第二维$2>1$，因此有$(1,2) > (1,1)$，因此就会返回区间$(1,2)$的迭代器，而这个迭代器就是st.begin()，如果再执行迭代器减$1$操作，就会越界报错。因此当我们要查找覆盖$x$的区间时，应该把第二维赋值为正无穷（这里下标最大不超过$n$，因此可以用$n+1$来表示正无穷），这样即使pair的第一维相等，也会强制往后比较，这样就避免了这种边界情况。

　　更新：另一种并查集的实现方法，灵感来源于题目Subarray With Elements Greater Than Varying Threshold

　　AC代码如下：

class Solution {
public:
    vector<int> fa;
    vector<long long> sum;

    int find(int x) {
        return fa[x] == x ? fa[x] : fa[x] = find(fa[x]);
    }

    vector<long long> maximumSegmentSum(vector<int>& nums, vector<int>& removeQueries) {
        int n = nums.size();
        for (int i = 0; i <= n ; i++) {
            fa.push_back(i);
            sum.push_back(0);
        }

        vector<long long> ans;
        multiset<long long> mst;
        mst.insert(0);
        for (int i = n - 1; i >= 0; i--) {
            ans.push_back(*mst.rbegin());
            int x = removeQueries[i];
            if (sum[find(x)]) mst.erase(mst.find(sum[find(x)]));
            if (x + 1 < n && sum[find(x + 1)]) mst.erase(mst.find(sum[find(x + 1)]));
            sum[find(x + 1)] += sum[x] + nums[x];
            fa[x] = find(x + 1);
            mst.insert(sum[find(x + 1)]);
        }

        reverse(ans.begin(), ans.end());
        return ans;
    }
};

　　反着做就是一开始没有任何数，然后每次都往数轴上加一个数，如果可以合并的话就合并成一个更大的区间，最后整个数组填满了数形成一个与原来一样的$1 \sim n$的区间。可以发现执行过程就是区间不断合并的过程（维护连通性），因此可以用并查集。合并后还需要求新区间的总和，然后用这个总和来更新最大值。

　　AC代码如下：

class Solution {
public:
    vector<int> fa;
    vector<long long> sum;

    int find(int x) {
        return fa[x] == x ? fa[x] : fa[x] = find(fa[x]);
    }

    vector<long long> maximumSegmentSum(vector<int>& nums, vector<int>& removeQueries) {
        int n = nums.size();
        for (int i = 0; i < n; i++) {
            fa.push_back(i);
            sum.push_back(nums[i]);
        }

        vector<bool> vis(n);    // vis[i] == false表示第i个数还没有被加进来
        vector<long long> ans;
        long long maxs = 0;
        for (int i = n - 1; i >= 0; i--) {
            ans.push_back(maxs);
            int x = removeQueries[i];
            vis[x] = true;  // x这个位置的数被加进来
            if (x - 1 >= 0 && vis[x - 1]) { // 合并x-1位置的连通块
                sum[x] += sum[find(x - 1)];
                fa[find(x - 1)] = find(x);
            }
            if (x + 1 < n && vis[x + 1]) {  // 合并x+1位置的连通块
                sum[x] += sum[find(x + 1)];
                fa[find(x + 1)] = x;
            }
            maxs = max(maxs, sum[x]);
        }
        
        reverse(ans.begin(), ans.end());
        return ans;
    }
};

 

参考资料

　　力扣第85场双周赛：https://www.bilibili.com/video/BV1u14y1t771