LeetCode 2926 - Maximum Balanced Subsequence Sum

This problem asks us to find the maximum sum of a balanced subsequence from a given integer array nums. A subsequence is a selection of elements from the array in their original order, possibly skipping elements.

leetcodehardarraybinary-searchdynamic-programmingbinary-indexed-treesegment-tree

LeetCode Problem 2926

Difficulty: 🔴 Hard
Topics: Array, Binary Search, Dynamic Programming, Binary Indexed Tree, Segment Tree

Solution

Problem Understanding

This problem asks us to find the maximum sum of a balanced subsequence from a given integer array nums. A subsequence is a selection of elements from the array in their original order, possibly skipping elements. A subsequence is balanced if, for every pair of consecutive indices i and j in the subsequence, the difference in values nums[j] - nums[i] is at least the difference in indices j - i. A subsequence of length 1 is automatically considered balanced.

The input array can be quite large, up to 10^5 elements, and numbers can range from -10^9 to 10^9. We must therefore find an efficient approach that does not attempt to enumerate all subsequences, which would be exponential in time. Edge cases include arrays with all negative numbers, arrays of length 1, and arrays where every pair of consecutive elements violates the balanced condition.

Approaches

The brute-force approach is to generate all possible subsequences, check whether each subsequence is balanced according to the definition, and compute their sums. While correct, this is infeasible because the number of subsequences grows exponentially as 2^n, which is far too slow for n = 10^5.

The key insight for an optimal approach comes from observing that the condition nums[j] - nums[i] >= j - i can be rewritten as nums[j] - j >= nums[i] - i. This transformation allows us to treat the problem as a variant of the maximum sum increasing subsequence problem on the array nums[i] - i. With this observation, we can use dynamic programming with a segment tree or binary indexed tree to maintain the maximum sum of subsequences ending at previous indices efficiently.

Approach Time Complexity Space Complexity Notes
Brute Force O(2^n) O(n) Enumerates all subsequences, checks balance, sums elements
Optimal O(n log n) O(n) Uses transformed values nums[i]-i and a segment tree or binary indexed tree to track maximum sums

Algorithm Walkthrough

  1. Transform the array by computing adjusted[i] = nums[i] - i. This simplifies the balanced condition to adjusted[j] >= adjusted[i] for valid subsequences.
  2. Coordinate compression: Since values of adjusted[i] may be large or negative, compress them into a smaller range to use with a segment tree or binary indexed tree efficiently.
  3. Initialize a segment tree or binary indexed tree to keep track of the maximum subsequence sum for any ending value seen so far.
  4. Iterate through the array from left to right:
  • For each i, query the segment tree for the maximum sum ending at any value <= adjusted[i].
  • Compute the potential new maximum sum including nums[i].
  • Update the segment tree at the compressed index corresponding to adjusted[i] with the new maximum sum.
  1. Return the overall maximum sum tracked during the iteration.

Why it works: Transforming nums[i] - i converts the index difference condition into a monotonic non-decreasing condition, allowing us to treat it like a maximum sum increasing subsequence problem. Using a segment tree ensures that we efficiently maintain and query the maximum sum for all feasible previous elements.

Python Solution

from typing import List
import bisect

class Solution:
    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:
        n = len(nums)
        adjusted = [nums[i] - i for i in range(n)]
        
        # Coordinate compression
        sorted_unique = sorted(set(adjusted))
        index_map = {val: idx for idx, val in enumerate(sorted_unique)}
        size = len(sorted_unique)
        
        # Segment tree for max query and update
        tree = [float('-inf')] * (2 * size)
        
        def update(idx: int, value: int):
            idx += size
            tree[idx] = max(tree[idx], value)
            while idx > 1:
                idx //= 2
                tree[idx] = max(tree[2*idx], tree[2*idx+1])
        
        def query(l: int, r: int) -> int:
            l += size
            r += size
            res = float('-inf')
            while l <= r:
                if l % 2 == 1:
                    res = max(res, tree[l])
                    l += 1
                if r % 2 == 0:
                    res = max(res, tree[r])
                    r -= 1
                l //= 2
                r //= 2
            return res
        
        max_sum = float('-inf')
        for i in range(n):
            idx = index_map[adjusted[i]]
            best_prev = query(0, idx)  # max sum ending at <= adjusted[i]
            current_sum = nums[i] + (best_prev if best_prev != float('-inf') else 0)
            update(idx, current_sum)
            max_sum = max(max_sum, current_sum)
        
        return max_sum

The Python implementation first transforms the array using nums[i] - i, compresses the values to reduce segment tree size, and maintains a segment tree to query and update the maximum subsequence sum. Each element is processed efficiently in O(log n) time due to the segment tree operations.

Go Solution

package main

import (
    "math"
    "sort"
)

func maxBalancedSubsequenceSum(nums []int) int64 {
    n := len(nums)
    adjusted := make([]int64, n)
    for i := 0; i < n; i++ {
        adjusted[i] = int64(nums[i]) - int64(i)
    }
    
    // Coordinate compression
    unique := make([]int64, n)
    copy(unique, adjusted)
    sort.Slice(unique, func(i, j int) bool { return unique[i] < unique[j] })
    unique = removeDuplicates(unique)
    
    indexMap := make(map[int64]int)
    for idx, val := range unique {
        indexMap[val] = idx
    }
    size := len(unique)
    
    tree := make([]int64, 2*size)
    for i := range tree {
        tree[i] = math.MinInt64
    }
    
    var update func(int, int64)
    update = func(idx int, value int64) {
        idx += size
        if value > tree[idx] {
            tree[idx] = value
        }
        for idx > 1 {
            idx /= 2
            tree[idx] = max(tree[2*idx], tree[2*idx+1])
        }
    }
    
    query := func(l, r int) int64 {
        res := math.MinInt64
        l += size
        r += size
        for l <= r {
            if l%2 == 1 {
                res = max(res, tree[l])
                l++
            }
            if r%2 == 0 {
                res = max(res, tree[r])
                r--
            }
            l /= 2
            r /= 2
        }
        return res
    }
    
    maxSum := int64(math.MinInt64)
    for i := 0; i < n; i++ {
        idx := indexMap[adjusted[i]]
        bestPrev := query(0, idx)
        currentSum := int64(nums[i])
        if bestPrev != math.MinInt64 {
            currentSum += bestPrev
        }
        update(idx, currentSum)
        if currentSum > maxSum {
            maxSum = currentSum
        }
    }
    
    return maxSum
}

func removeDuplicates(arr []int64) []int64 {
    n := len(arr)
    if n == 0 {
        return arr
    }
    res := []int64{arr[0]}
    for i := 1; i < n; i++ {
        if arr[i] != arr[i-1] {
            res = append(res, arr[i])
        }
    }
    return res
}

func max(a, b int64) int64 {
    if a > b {
        return a
    }
    return b
}

The Go solution mirrors the Python logic. Differences include explicit handling of int64 to avoid integer overflow and a helper function to remove duplicates during coordinate compression.

Worked Examples

Example 1

Input: nums = [3,3,5,6]

i nums[i] adjusted[i] best_prev current_sum tree update
0 3 3-0=3 -inf 3 update(3,3)
1 3 3-1=2 -inf 3 update(2,3)
2 5 5-2=3 3 8 update(3,8)
3 6 6-3=3 8 14 update(3,14)

Output: 14

Example 2

Input: nums = [5,-1,-3,8]

| i | nums[i] | adjusted[i] | best_prev | current_sum | tree update |

|---