Maximum Balanced Subsequence Sum

Problem

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 themaximum 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.

Examples

Example 1

    
    
    Input: nums = [3,3,5,6]
    Output: 14
    Explanation: In this example, the subsequence [3,5,6] consisting of indices 0, 2, and 3 can be selected.
    nums[2] - nums[0] >= 2 - 0.
    nums[3] - nums[2] >= 3 - 2.
    Hence, it is a balanced subsequence, and its sum is the maximum among the balanced subsequences of nums.
    The subsequence consisting of indices 1, 2, and 3 is also valid.
    It can be shown that it is not possible to get a balanced subsequence with a sum greater than 14.

Example 2

    
    
    Input: nums = [5,-1,-3,8]
    Output: 13
    Explanation: In this example, the subsequence [5,8] consisting of indices 0 and 3 can be selected.
    nums[3] - nums[0] >= 3 - 0.
    Hence, it is a balanced subsequence, and its sum is the maximum among the balanced subsequences of nums.
    It can be shown that it is not possible to get a balanced subsequence with a sum greater than 13.
    

Example 3

    
    
    Input: nums = [-2,-1]
    Output: -1
    Explanation: In this example, the subsequence [-1] can be selected.
    It is a balanced subsequence, and its sum is the maximum among the balanced subsequences of nums.
    

Constraints

• 1 <= nums.length <= 10^5
• -10^9 <= nums[i] <= 10^9

Solution

Method 1 – Dynamic Programming with Coordinate Compression and Fenwick Tree

Intuition

To maximize the sum of a balanced subsequence, we want to extend the subsequence as much as possible. For each index, we can use dynamic programming to track the best sum ending at each position, and use a Fenwick Tree (Binary Indexed Tree) to efficiently query and update the best sum for valid previous indices. Since the difference nums[j] - j must be non-decreasing, we use coordinate compression on nums[i] - i.

Approach

1. For each index i, compute key = nums[i] - i and compress all such keys.
2. Use a Fenwick Tree to maintain the maximum sum for each compressed key.
3. For each index i in order:
   • Query the Fenwick Tree for the maximum sum for keys ≤ current key.
   • Set dp[i] = nums[i] + max_sum.
   • Update the Fenwick Tree at the current key with dp[i].
4. The answer is the maximum value in dp.

Code

C++

class Fenwick {
public:
    vector<long long> tree;
    int n;
    Fenwick(int n): n(n), tree(n+2, LLONG_MIN) {}
    void update(int i, long long val) {
        for (++i; i <= n+1; i += i&-i) tree[i] = max(tree[i], val);
    }
    long long query(int i) {
        long long res = LLONG_MIN;
        for (++i; i > 0; i -= i&-i) res = max(res, tree[i]);
        return res;
    }
};
class Solution {
public:
    long long maxBalancedSubsequenceSum(vector<int>& nums) {
        int n = nums.size();
        vector<long long> keys;
        for (int i = 0; i < n; ++i) keys.push_back((long long)nums[i] - i);
        sort(keys.begin(), keys.end());
        keys.erase(unique(keys.begin(), keys.end()), keys.end());
        Fenwick bit(keys.size());
        long long ans = LLONG_MIN;
        for (int i = 0; i < n; ++i) {
            int idx = lower_bound(keys.begin(), keys.end(), (long long)nums[i] - i) - keys.begin();
            long long best = bit.query(idx);
            long long cur = nums[i];
            if (best != LLONG_MIN) cur += best;
            bit.update(idx, cur);
            ans = max(ans, cur);
        }
        return ans;
    }
};

Go

func maxBalancedSubsequenceSum(nums []int) int64 {
    n := len(nums)
    keys := make([]int, n)
    for i := 0; i < n; i++ {
        keys[i] = nums[i] - i
    }
    uniq := make([]int, len(keys))
    copy(uniq, keys)
    sort.Ints(uniq)
    uniq = unique(uniq)
    idx := func(x int) int {
        return sort.SearchInts(uniq, x)
    }
    bit := make([]int64, len(uniq)+2)
    for i := range bit {
        bit[i] = math.MinInt64
    }
    ans := int64(math.MinInt64)
    for i := 0; i < n; i++ {
        k := idx(nums[i]-i)
        best := int64(math.MinInt64)
        for j := k+1; j > 0; j -= j&-j {
            if bit[j] > best {
                best = bit[j]
            }
        }
        cur := int64(nums[i])
        if best != math.MinInt64 {
            cur += best
        }
        for j := k+1; j < len(bit); j += j&-j {
            if bit[j] < cur {
                bit[j] = cur
            }
        }
        if cur > ans {
            ans = cur
        }
    }
    return ans
}
func unique(a []int) []int {
    if len(a) == 0 { return a }
    res := []int{a[0]}
    for i := 1; i < len(a); i++ {
        if a[i] != a[i-1] {
            res = append(res, a[i])
        }
    }
    return res
}

Java

class Solution {
    public long maxBalancedSubsequenceSum(int[] nums) {
        int n = nums.length;
        long[] keys = new long[n];
        for (int i = 0; i < n; ++i) keys[i] = (long)nums[i] - i;
        long[] uniq = Arrays.stream(keys).distinct().sorted().toArray();
        Map<Long, Integer> idx = new HashMap<>();
        for (int i = 0; i < uniq.length; ++i) idx.put(uniq[i], i);
        long[] bit = new long[uniq.length+2];
        Arrays.fill(bit, Long.MIN_VALUE);
        long ans = Long.MIN_VALUE;
        for (int i = 0; i < n; ++i) {
            int k = idx.get((long)nums[i] - i);
            long best = Long.MIN_VALUE;
            for (int j = k+1; j > 0; j -= j&-j) best = Math.max(best, bit[j]);
            long cur = nums[i];
            if (best != Long.MIN_VALUE) cur += best;
            for (int j = k+1; j < bit.length; j += j&-j) bit[j] = Math.max(bit[j], cur);
            ans = Math.max(ans, cur);
        }
        return ans;
    }
}

Kotlin

class Solution {
    fun maxBalancedSubsequenceSum(nums: IntArray): Long {
        val n = nums.size
        val keys = LongArray(n) { nums[it].toLong() - it }
        val uniq = keys.distinct().sorted()
        val idx = uniq.withIndex().associate { it.value to it.index }
        val bit = LongArray(uniq.size+2) { Long.MIN_VALUE }
        var ans = Long.MIN_VALUE
        for (i in 0 until n) {
            val k = idx[nums[i].toLong() - i]!!
            var best = Long.MIN_VALUE
            var j = k+1
            while (j > 0) {
                best = maxOf(best, bit[j])
                j -= j and -j
            }
            var cur = nums[i].toLong()
            if (best != Long.MIN_VALUE) cur += best
            j = k+1
            while (j < bit.size) {
                bit[j] = maxOf(bit[j], cur)
                j += j and -j
            }
            ans = maxOf(ans, cur)
        }
        return ans
    }
}

Python

class Solution:
    def maxBalancedSubsequenceSum(self, nums: list[int]) -> int:
        n = len(nums)
        keys = [nums[i] - i for i in range(n)]
        uniq = sorted(set(keys))
        idx = {v: i for i, v in enumerate(uniq)}
        bit = [float('-inf')] * (len(uniq)+2)
        def query(i):
            res = float('-inf')
            i += 1
            while i > 0:
                res = max(res, bit[i])
                i -= i & -i
            return res
        def update(i, val):
            i += 1
            while i < len(bit):
                bit[i] = max(bit[i], val)
                i += i & -i
        ans = float('-inf')
        for i in range(n):
            k = idx[nums[i] - i]
            best = query(k)
            cur = nums[i]
            if best != float('-inf'):
                cur += best
            update(k, cur)
            ans = max(ans, cur)
        return ans

Rust

impl Solution {
    pub fn max_balanced_subsequence_sum(nums: Vec<i32>) -> i64 {
        let n = nums.len();
        let mut keys: Vec<i64> = (0..n).map(|i| nums[i] as i64 - i as i64).collect();
        keys.sort();
        keys.dedup();
        let mut bit = vec![i64::MIN; keys.len()+2];
        let mut ans = i64::MIN;
        for (i, &x) in nums.iter().enumerate() {
            let k = keys.binary_search(&(x as i64 - i as i64)).unwrap();
            let mut best = i64::MIN;
            let mut j = k+1;
            while j > 0 {
                best = best.max(bit[j]);
                j -= j & (!j + 1);
            }
            let mut cur = x as i64;
            if best != i64::MIN {
                cur += best;
            }
            j = k+1;
            while j < bit.len() {
                bit[j] = bit[j].max(cur);
                j += j & (!j + 1);
            }
            ans = ans.max(cur);
        }
        ans
    }
}

TypeScript

class Solution {
    maxBalancedSubsequenceSum(nums: number[]): number {
        const n = nums.length;
        const keys = nums.map((x, i) => x - i);
        const uniq = Array.from(new Set(keys)).sort((a, b) => a - b);
        const idx = new Map(uniq.map((v, i) => [v, i]));
        const bit = Array(uniq.length+2).fill(-Infinity);
        function query(i: number): number {
            let res = -Infinity;
            i++;
            while (i > 0) {
                res = Math.max(res, bit[i]);
                i -= i & -i;
            }
            return res;
        }
        function update(i: number, val: number) {
            i++;
            while (i < bit.length) {
                bit[i] = Math.max(bit[i], val);
                i += i & -i;
            }
        }
        let ans = -Infinity;
        for (let i = 0; i < n; ++i) {
            const k = idx.get(nums[i] - i)!;
            const best = query(k);
            let cur = nums[i];
            if (best !== -Infinity) cur += best;
            update(k, cur);
            ans = Math.max(ans, cur);
        }
        return ans;
    }
}

Complexity

• ⏰ Time complexity: O(n log n), due to coordinate compression and Fenwick Tree operations.
• 🧺 Space complexity: O(n), for the DP and Fenwick Tree arrays.