We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

Maximum Balanced Subsequence Sum

Number: 3184

Difficulty: Hard

Paid? No

Companies: Amazon


Problem Description

Given an array of integers nums, find a balanced subsequence with the maximum possible sum. A subsequence is balanced if for every two consecutive indices in the subsequence, say i and j (with i < j), the condition nums[j] - nums[i] >= j - i holds. Note that any single element is a balanced subsequence by itself.


Key Insights

  • The balanced condition, nums[j] - nums[i] >= j - i, can be rearranged as (nums[j] - j) >= (nums[i] - i). Thus, for any valid subsequence (taken in the original order), the sequence (nums[i] - i) must be non-decreasing.
  • This formulation transforms the problem into finding a maximum-sum subsequence subject to a non-decreasing constraint on (nums[i] - i).
  • We can adopt a dynamic programming approach where for each index j we compute the maximum sum of a balanced subsequence ending at j by considering all previous indices i with (nums[i]-i) <= (nums[j]-j).
  • To efficiently retrieve the best previous DP state, a Binary Indexed Tree (BIT) or Segment Tree can be leveraged after coordinate-compressing the possible values of (nums[i]-i).

Space and Time Complexity

Time Complexity: O(n · log n), where n is the length of the nums array (BIT queries and updates happen in logarithmic time after compression)
Space Complexity: O(n), for storing the DP array and BIT data structure


Solution

The idea is to use dynamic programming along with a Binary Indexed Tree (BIT) to quickly access the optimum (maximum sum) from previous indices that satisfy our condition. For each index j, define dp[j] as the maximum sum of a balanced subsequence ending at index j. Then:

dp[j] = nums[j] + max{ dp[i] } over all i < j such that (nums[i] - i) <= (nums[j] - j)

If there is no such index i, then dp[j] is simply nums[j].

Since the values (nums[i]-i) vary widely and can be negative, we first perform coordinate compression. As we process the array from left-to-right, we update our BIT at the index corresponding to (nums[j]-j) with dp[j]. BIT queries give the running maximum dp for all positions with compressed value less than or equal to the current.

This algorithm ensures that the overall final answer is the maximum dp[j] over all indices j.


Code Solutions

# Python implementation using BIT (Fenwick Tree)

# Class for BIT which supports point update and range maximum query.
class BIT:
    def __init__(self, size):
        self.size = size
        self.tree = [-10**18] * (size + 1)  # initialize with very small numbers
        
    # Update the BIT at index i (1-indexed) with value 'val'
    def update(self, i, val):
        while i <= self.size:
            self.tree[i] = max(self.tree[i], val)
            i += i & -i
            
    # Query the maximum value in the range [1, i]
    def query(self, i):
        result = -10**18
        while i > 0:
            result = max(result, self.tree[i])
            i -= i & -i
        return result

def maximumBalancedSubsequenceSum(nums):
    n = len(nums)
    # Calculate keys: key[i] = nums[i] - i
    keys = [nums[i] - i for i in range(n)]
    
    # Coordinate compression for keys
    sorted_keys = sorted(set(keys))
    comp = {val: idx+1 for idx, val in enumerate(sorted_keys)}  # +1 to make BIT 1-indexed
    
    # Initialize BIT with size equal to number of unique keys
    bit = BIT(len(sorted_keys))
    
    max_sum = -10**18  # track global maximum
    
    # Process each index in nums
    for i in range(n):
        key = comp[keys[i]]  # get compressed key for nums[i]-i
        # Query BIT for maximum dp value for any j with keys[j] <= current key
        best_sum = bit.query(key)
        if best_sum == -10**18:
            dp = nums[i]  # no previous valid index, start a new subsequence
        else:
            dp = best_sum + nums[i]
        # Update BIT with the new dp at the current key position
        bit.update(key, dp)
        max_sum = max(max_sum, dp)
    
    return max_sum

# Example run:
print(maximumBalancedSubsequenceSum([3, 3, 5, 6]))  # Expected output: 14
← Back to All Questions