Leetcode 2926. Maximum Balanced Subsequence Sum

You are given a 0-indexed integer array nums.

A subsequence of nums having length k and consisting of indices i0 < i1 < … < ik-1 is balanced if the following holds:

nums[ij] - nums[ij-1] >= ij - ij-1, for every j in the range [1, k - 1].

A subsequence of nums having length 1 is considered balanced.

Return an integer denoting the maximum possible sum of elements in a balanced subsequence of nums.

A subsequence of an array is a new non-empty array that is formed from the original array by deleting some (possibly none) of the elements without disturbing the relative positions of the remaining elements.

Firstly, we can imagine that when all the elements are negative we can directly return the maximum negatives.

Next, we need to find the increasing subsequence for $\text{nums[ij]}-ij$ with the maximum sum of $\text{nums[ij]}$. Our idea is to build a segment tree, which stores the maximum sum when $ij$ is the index of the ending element of the subsequence. We search the elements in the ascending order of $\text{nums[ij]}-ij$ so that we can guarantee that all the updated elements in the segment tree satisfies the increasing condition to become the precedent elements. Then we search the range of $[0,ij]$ in the segment tree for the update. The time complexity is $\mathcal{O}(N\log (N))$

class Solution {
public:
    void build(vector<int>& nums, vector<int>& tree, int v, int l, int r) {
        if (l == r) tree[v] = nums[l];
        else {
            int m = (l + r) / 2;
            build(nums, tree, v*2+1, l,   m);
            build(nums, tree, v*2+2, m+1, r);
            tree[v] = min(tree[v*2+1], tree[v*2+2]);
        }
    }
    long long combine(long long a, long long b) {
        // if (a < 0 && b < 0) {
        //     return max(a, b);
        // }
        return max(a, b);
    }

    void update(vector<long long>& tree, int n, int ind, long long x) {
        int v = 0, l = 0, r = n-1;
        while (l != r) {
            int m = (l + r) / 2;
            if (ind <= m) {
                v = v*2+1;
                r = m;
            } else {
                v = v*2+2;
                l = m+1;
            }
        }
        tree[v] = x;
        while (v!=0) {
            int u = (v+1)/2-1;
            tree[u] = combine(tree[u*2+1], tree[u*2+2]);
            v = u;
        }
    }

    long long query(vector<long long>& tree, int v, int tl, int tr, int l, int r) {
        if (l > r) return 0;
        if (tl == l && tr == r) return tree[v];
        int tm = (tl + tr) / 2;
        return combine(
            query(tree, v*2+1, tl, tm, l, min(tm, r)), 
            query(tree, v*2+2, tm+1, tr, max(l, tm+1), r)
        );
    }

    long long maxBalancedSubsequenceSum(vector<int>& nums) {
        long long mx = *max_element(nums.begin(), nums.end());
        if (mx <= 0) return mx;

        vector<pair<int, int> > balance;
        vector<long long> tree(nums.size()*4, 0);

        for (int i = 0; i < nums.size(); i++) {
            balance.push_back(make_pair(nums[i]-i, i));
        }
        sort(balance.begin(), balance.end());

        long long ans = 0;
        for (auto it = balance.begin(); it != balance.end(); ++it) {
            long long cur;
            cur = query(tree, 0, 0, nums.size()-1, 0, it->second) + nums[it->second];

            update(tree, nums.size(), it->second, cur);
            ans = max(ans, cur);
        }
        return ans;
    }
};