2926. Maximum Balanced Subsequence Sum

HardBinary Indexed TreeSegment TreeArrayBinary SearchDynamic Programming

Problem Description

You are given a 0-indexed integer array nums.

A subsequence of nums with length k and consisting of indices i₀ < i₁ < ... < i_{k-1} is called balanced if it satisfies the following condition:

For every j in the range [1, k-1]: nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1}

In other words, the difference between consecutive values in the subsequence must be at least as large as the difference between their indices.

A subsequence of length 1 is automatically considered balanced.

Your task is to find the maximum possible sum of elements in any balanced subsequence of nums.

Key Points:

A subsequence maintains the relative order of elements from the original array but doesn't need to be contiguous
The balance condition requires that the value increase between consecutive elements is at least as much as the index increase
You need to maximize the sum of the selected elements while maintaining the balance property

Example Understanding: If you select indices 2 and 5 from the array, then:

The values are nums[2] and nums[5]
The balance condition requires: nums[5] - nums[2] >= 5 - 2 = 3
These two indices can form a balanced subsequence only if the value at index 5 is at least 3 more than the value at index 2

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

Let's start by understanding what the balance condition really means. We have nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1} for consecutive elements in our subsequence.

If we rearrange this inequality, we get: nums[i_j] - i_j >= nums[i_{j-1}] - i_{j-1}

This is a key insight! If we create a new array arr where arr[i] = nums[i] - i, then the balance condition simply becomes: for any two consecutive indices in our subsequence, the corresponding values in arr must be non-decreasing.

In other words, a balanced subsequence in the original array corresponds to selecting indices where the arr values form a non-decreasing sequence.

Now our problem transforms into: given the array arr, select a subset of indices such that:

The arr values at these indices are non-decreasing
The sum of the corresponding nums values is maximized

This looks like a dynamic programming problem where for each index i, we need to find the best subsequence ending at i. We can define f[i] as the maximum sum achievable by a balanced subsequence ending at index i.

For each position i, we look at all previous positions j where arr[j] <= arr[i] (meaning we can extend the subsequence from j to i), and choose the one that gives us the maximum sum. The recurrence is: f[i] = max(f[j] for all j < i where arr[j] <= arr[i]) + nums[i]

However, finding all such j values for each i would be O(n²). To optimize this, we need a data structure that can efficiently:

Query the maximum f[j] value among all j where arr[j] <= arr[i]
Update the value at a given position

A Binary Indexed Tree (Fenwick Tree) is perfect for this! We can use coordinate compression to map the arr values to a smaller range, then use the BIT to maintain prefix maximums. This reduces our complexity to O(n log n).

Learn more about Segment Tree, Binary Search and Dynamic Programming patterns.

Solution Implementation

1from typing import List 2from bisect import bisect_left 3 4class BinaryIndexedTree: 5 """ 6 Binary Indexed Tree (Fenwick Tree) for range maximum queries. 7 This implementation supports point updates and prefix maximum queries. 8 """ 9 10 def __init__(self, n: int): 11 """ 12 Initialize the Binary Indexed Tree with size n. 13 14 Args: 15 n: The size of the tree (1-indexed) 16 """ 17 self.size = n 18 # Tree array initialized with negative infinity (1-indexed) 19 self.tree = [float('-inf')] * (n + 1) 20 21 def update(self, index: int, value: int): 22 """ 23 Update the tree at position index with maximum of current and new value. 24 25 Args: 26 index: The position to update (1-indexed) 27 value: The new value to consider for maximum 28 """ 29 while index <= self.size: 30 # Update current position with maximum value 31 self.tree[index] = max(self.tree[index], value) 32 # Move to next position affected by this update 33 # Adding the least significant bit moves to parent in update tree 34 index += index & -index 35 36 def query(self, index: int) -> int: 37 """ 38 Query the maximum value in range [1, index]. 39 40 Args: 41 index: The right boundary of the query range (1-indexed) 42 43 Returns: 44 The maximum value in the range [1, index] 45 """ 46 maximum_value = float('-inf') 47 while index > 0: 48 # Consider current position's value 49 maximum_value = max(maximum_value, self.tree[index]) 50 # Move to previous range by removing least significant bit 51 index -= index & -index 52 return maximum_value 53 54 55class Solution: 56 def maxBalancedSubsequenceSum(self, nums: List[int]) -> int: 57 """ 58 Find the maximum sum of a balanced subsequence. 59 A subsequence is balanced if for any two consecutive elements at positions i and j, 60 nums[j] - nums[i] >= j - i. 61 62 Args: 63 nums: The input array of integers 64 65 Returns: 66 The maximum sum of a balanced subsequence 67 """ 68 # Transform the constraint: nums[j] - nums[i] >= j - i 69 # Rearranging: nums[j] - j >= nums[i] - i 70 # Create array of transformed values 71 transformed_values = [nums[i] - i for i in range(len(nums))] 72 73 # Compress coordinates: map transformed values to ranks 74 # This allows us to use them as indices in the BIT 75 sorted_unique_values = sorted(set(transformed_values)) 76 77 # Initialize Binary Indexed Tree with compressed coordinate space size 78 bit = BinaryIndexedTree(len(sorted_unique_values)) 79 80 # Process each element in the original array 81 for i, current_num in enumerate(nums): 82 # Find the compressed coordinate (rank) of the transformed value 83 # Add 1 because BIT is 1-indexed 84 compressed_index = bisect_left(sorted_unique_values, current_num - i) + 1 85 86 # Query maximum sum ending at or before current transformed value 87 # Take max with 0 to allow starting a new subsequence 88 max_previous_sum = max(bit.query(compressed_index), 0) 89 90 # Calculate new sum including current element 91 new_sum = max_previous_sum + current_num 92 93 # Update the BIT with the new maximum sum at this position 94 bit.update(compressed_index, new_sum) 95 96 # Return the overall maximum sum found 97 return bit.query(len(sorted_unique_values)) 98

1/** 2 * Binary Indexed Tree (Fenwick Tree) for range maximum queries 3 * This implementation maintains maximum values instead of sums 4 */ 5class BinaryIndexedTree { 6 private int size; 7 private long[] tree; 8 private final long NEGATIVE_INFINITY = 1L << 60; 9 10 /** 11 * Initialize BIT with given size 12 * @param size the size of the tree 13 */ 14 public BinaryIndexedTree(int size) { 15 this.size = size; 16 this.tree = new long[size + 1]; 17 Arrays.fill(this.tree, -NEGATIVE_INFINITY); 18 } 19 20 /** 21 * Update the tree with maximum value at position x 22 * @param position the 1-indexed position to update 23 * @param value the value to potentially update with (takes max) 24 */ 25 public void update(int position, long value) { 26 while (position <= size) { 27 tree[position] = Math.max(tree[position], value); 28 // Move to next position affected by this update 29 position += position & -position; 30 } 31 } 32 33 /** 34 * Query the maximum value from index 1 to x 35 * @param position the 1-indexed position to query up to 36 * @return the maximum value in range [1, position] 37 */ 38 public long query(int position) { 39 long maximum = -NEGATIVE_INFINITY; 40 while (position > 0) { 41 maximum = Math.max(maximum, tree[position]); 42 // Move to parent node 43 position -= position & -position; 44 } 45 return maximum; 46 } 47} 48 49/** 50 * Solution for finding maximum balanced subsequence sum 51 * A subsequence is balanced if nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1} 52 * This can be rewritten as nums[i_j] - i_j >= nums[i_{j-1}] - i_{j-1} 53 */ 54class Solution { 55 /** 56 * Find the maximum sum of a balanced subsequence 57 * @param nums the input array 58 * @return the maximum sum of any balanced subsequence 59 */ 60 public long maxBalancedSubsequenceSum(int[] nums) { 61 int n = nums.length; 62 63 // Transform array: store nums[i] - i for each position 64 // This transformation helps identify valid balanced subsequences 65 int[] transformedValues = new int[n]; 66 for (int i = 0; i < n; ++i) { 67 transformedValues[i] = nums[i] - i; 68 } 69 70 // Sort and remove duplicates to create coordinate compression 71 Arrays.sort(transformedValues); 72 int uniqueCount = 0; 73 for (int i = 0; i < n; ++i) { 74 if (i == 0 || transformedValues[i] != transformedValues[i - 1]) { 75 transformedValues[uniqueCount++] = transformedValues[i]; 76 } 77 } 78 79 // Initialize BIT with compressed coordinate space 80 BinaryIndexedTree fenwickTree = new BinaryIndexedTree(uniqueCount); 81 82 // Process each element in original order 83 for (int i = 0; i < n; ++i) { 84 // Find compressed coordinate of current transformed value 85 int compressedIndex = binarySearch(transformedValues, nums[i] - i, uniqueCount) + 1; 86 87 // Query maximum sum ending at values <= current transformed value 88 // Add current value to extend the subsequence 89 long maxSum = Math.max(fenwickTree.query(compressedIndex), 0) + nums[i]; 90 91 // Update the tree with new maximum at this position 92 fenwickTree.update(compressedIndex, maxSum); 93 } 94 95 // Return the overall maximum from all possible subsequences 96 return fenwickTree.query(uniqueCount); 97 } 98 99 /** 100 * Binary search to find leftmost position of target value 101 * @param array the sorted array to search 102 * @param target the value to search for 103 * @param right the right boundary (exclusive) 104 * @return the index of first occurrence of target or where it would be inserted 105 */ 106 private int binarySearch(int[] array, int target, int right) { 107 int left = 0; 108 while (left < right) { 109 int mid = (left + right) >> 1; 110 if (array[mid] >= target) { 111 right = mid; 112 } else { 113 left = mid + 1; 114 } 115 } 116 return left; 117 } 118} 119

1class BinaryIndexedTree { 2private: 3 int size; 4 vector<long long> tree; 5 const long long NEGATIVE_INFINITY = 1e18; 6 7public: 8 // Initialize BIT with given size, all values set to negative infinity 9 BinaryIndexedTree(int n) { 10 this->size = n; 11 tree.resize(n + 1, -NEGATIVE_INFINITY); 12 } 13 14 // Update the BIT at position x with maximum value v 15 // Uses the BIT update pattern: move to next index by adding lowbit 16 void update(int index, long long value) { 17 while (index <= size) { 18 tree[index] = max(tree[index], value); 19 index += index & -index; // Add lowbit to move to next position 20 } 21 } 22 23 // Query the maximum value from index 1 to x 24 // Uses the BIT query pattern: move to previous index by subtracting lowbit 25 long long query(int index) { 26 long long maxValue = -NEGATIVE_INFINITY; 27 while (index > 0) { 28 maxValue = max(maxValue, tree[index]); 29 index -= index & -index; // Subtract lowbit to move to previous position 30 } 31 return maxValue; 32 } 33}; 34 35class Solution { 36public: 37 long long maxBalancedSubsequenceSum(vector<int>& nums) { 38 int n = nums.size(); 39 40 // Transform nums[i] to nums[i] - i to handle the balanced condition 41 // A balanced subsequence satisfies: nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1} 42 // Which transforms to: nums[i_j] - i_j >= nums[i_{j-1}] - i_{j-1} 43 vector<int> transformedValues(n); 44 for (int i = 0; i < n; ++i) { 45 transformedValues[i] = nums[i] - i; 46 } 47 48 // Coordinate compression: sort and remove duplicates 49 // This maps the transformed values to a smaller range [1, m] 50 sort(transformedValues.begin(), transformedValues.end()); 51 transformedValues.erase(unique(transformedValues.begin(), transformedValues.end()), 52 transformedValues.end()); 53 int compressedSize = transformedValues.size(); 54 55 // Initialize BIT with compressed size 56 BinaryIndexedTree fenwickTree(compressedSize); 57 58 // Process each element in original order 59 for (int i = 0; i < n; ++i) { 60 // Find the compressed index for current transformed value 61 // Add 1 because BIT uses 1-based indexing 62 int compressedIndex = lower_bound(transformedValues.begin(), 63 transformedValues.end(), 64 nums[i] - i) - transformedValues.begin() + 1; 65 66 // Calculate maximum sum ending at current position 67 // Either start new subsequence (0) or extend previous subsequence 68 long long currentMaxSum = max(fenwickTree.query(compressedIndex), 0LL) + nums[i]; 69 70 // Update BIT with the new maximum sum at this compressed index 71 fenwickTree.update(compressedIndex, currentMaxSum); 72 } 73 74 // Return the overall maximum sum by querying all positions 75 return fenwickTree.query(compressedSize); 76 } 77}; 78

1// Binary Indexed Tree (Fenwick Tree) for range maximum queries 2let n: number; 3let c: number[]; 4 5// Initialize the BIT with size and fill with negative infinity 6function initBIT(size: number): void { 7 n = size; 8 c = Array(size + 1).fill(-Infinity); 9} 10 11// Update the BIT at position x with maximum value v 12function update(x: number, v: number): void { 13 while (x <= n) { 14 c[x] = Math.max(c[x], v); 15 x += x & -x; // Move to next position using lowbit 16 } 17} 18 19// Query the maximum value from position 1 to x 20function query(x: number): number { 21 let maxValue = -Infinity; 22 while (x > 0) { 23 maxValue = Math.max(maxValue, c[x]); 24 x -= x & -x; // Move to parent position using lowbit 25 } 26 return maxValue; 27} 28 29// Binary search to find the leftmost position where arr[mid] >= target 30function binarySearch(arr: number[], target: number): number { 31 let left = 0; 32 let right = arr.length; 33 34 while (left < right) { 35 const mid = (left + right) >> 1; // Equivalent to Math.floor((left + right) / 2) 36 if (arr[mid] >= target) { 37 right = mid; 38 } else { 39 left = mid + 1; 40 } 41 } 42 return left; 43} 44 45// Find the maximum sum of a balanced subsequence 46// A subsequence is balanced if nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1} 47// Which can be rewritten as nums[i_j] - i_j >= nums[i_{j-1}] - i_{j-1} 48function maxBalancedSubsequenceSum(nums: number[]): number { 49 const length = nums.length; 50 51 // Create array of (nums[i] - i) values for comparison 52 const differences = Array(length).fill(0); 53 for (let i = 0; i < length; ++i) { 54 differences[i] = nums[i] - i; 55 } 56 57 // Sort differences to prepare for coordinate compression 58 differences.sort((a, b) => a - b); 59 60 // Remove duplicates from sorted array (coordinate compression) 61 let uniqueCount = 0; 62 for (let i = 0; i < length; ++i) { 63 if (i === 0 || differences[i] !== differences[i - 1]) { 64 differences[uniqueCount++] = differences[i]; 65 } 66 } 67 differences.length = uniqueCount; 68 69 // Initialize BIT with compressed size 70 initBIT(uniqueCount); 71 72 // Process each element in original order 73 for (let i = 0; i < length; ++i) { 74 // Find compressed index for current element's (nums[i] - i) value 75 const compressedIndex = binarySearch(differences, nums[i] - i) + 1; // +1 for 1-indexed BIT 76 77 // Query maximum sum ending before current position, add current value 78 const maxSum = Math.max(query(compressedIndex), 0) + nums[i]; 79 80 // Update BIT with new maximum sum at this position 81 update(compressedIndex, maxSum); 82 } 83 84 // Return the overall maximum sum 85 return query(uniqueCount); 86} 87

Solution Approach

Based on our intuition, let's implement the solution step by step:

Step 1: Transform the Problem

First, we create the transformed array arr where arr[i] = nums[i] - i:

arr = [x - i for i, x in enumerate(nums)]

Step 2: Coordinate Compression

Since we'll use a Binary Indexed Tree indexed by the values in arr, we need to compress these values to a smaller range [1, m] where m is the number of unique values:

s = sorted(set(arr)) # Get unique values and sort them

This creates a mapping where each unique value in arr gets an index from 1 to len(s).

Step 3: Binary Indexed Tree for Maximum Queries

We implement a BIT that supports:

update(x, v): Update position x with value v (keeping the maximum)
query(x): Get the maximum value in the range [1, x]

The BIT is initialized with negative infinity values since we're looking for maximums:

self.c = [-inf] * (n + 1)

Step 4: Dynamic Programming with BIT

For each element in nums, we:

Find the compressed index of arr[i] = nums[i] - i:

j = bisect_left(s, x - i) + 1 # +1 because BIT is 1-indexed

Query the maximum sum of all subsequences whose last element has arr value ≤ current arr[i]:

v = max(tree.query(j), 0) + x

We take max(..., 0) because we can start a new subsequence from the current element.

Update the BIT with the maximum sum ending at current position:

tree.update(j, v)

Step 5: Get the Final Answer

After processing all elements, the answer is the maximum value stored in the entire BIT:

return tree.query(len(s))

Time Complexity Analysis:

Sorting unique values: O(n log n)
For each element, we do binary search and BIT operations: O(n log n)
Overall: O(n log n)

Space Complexity: O(n) for the BIT and auxiliary arrays.

The key insight is that by transforming the problem and using coordinate compression with a BIT, we efficiently maintain the maximum sums for all valid prefixes, allowing us to build the optimal balanced subsequence incrementally.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a concrete example with nums = [3, 3, 5, 6].

Step 1: Transform the array We create arr[i] = nums[i] - i:

arr[0] = 3 - 0 = 3
arr[1] = 3 - 1 = 2
arr[2] = 5 - 2 = 3
arr[3] = 6 - 3 = 3

So arr = [3, 2, 3, 3]

Step 2: Coordinate Compression Get unique sorted values: s = [2, 3] This maps:

Value 2 → compressed index 1
Value 3 → compressed index 2

Step 3: Initialize BIT Create a BIT of size 3 (len(s) + 1) initialized with -inf

Step 4: Process each element

Processing index 0 (nums[0] = 3):

arr[0] = 3 → compressed index = 2
Query BIT for max in range [1, 2]: returns -inf
Calculate: v = max(-inf, 0) + 3 = 0 + 3 = 3
Update BIT at index 2 with value 3
BIT state: [-inf, -inf, 3]

Processing index 1 (nums[1] = 3):

arr[1] = 2 → compressed index = 1
Query BIT for max in range [1, 1]: returns -inf
Calculate: v = max(-inf, 0) + 3 = 0 + 3 = 3
Update BIT at index 1 with value 3
BIT state: [-inf, 3, 3]

Processing index 2 (nums[2] = 5):

arr[2] = 3 → compressed index = 2
Query BIT for max in range [1, 2]: returns max(3, 3) = 3
Calculate: v = max(3, 0) + 5 = 3 + 5 = 8
Update BIT at index 2 with value 8
BIT state: [-inf, 3, 8]

Processing index 3 (nums[3] = 6):

arr[3] = 3 → compressed index = 2
Query BIT for max in range [1, 2]: returns max(3, 8) = 8
Calculate: v = max(8, 0) + 6 = 8 + 6 = 14
Update BIT at index 2 with value 14
BIT state: [-inf, 3, 14]

Step 5: Get the answer Query the entire BIT: tree.query(2) = 14

Verification: The optimal subsequence is indices [1, 2, 3] with values [3, 5, 6]:

Check balance between indices 1 and 2: 5 - 3 = 2 >= 2 - 1 = 1 ✓
Check balance between indices 2 and 3: 6 - 5 = 1 >= 3 - 2 = 1 ✓
Sum = 3 + 5 + 6 = 14 ✓

This subsequence works because in the transformed array, we have arr[1] = 2 ≤ arr[2] = 3 ≤ arr[3] = 3, forming a non-decreasing sequence.

Time and Space Complexity

Time Complexity: O(n × log n)

The time complexity is dominated by several operations:

Creating the array arr with x - i for each element: O(n)
Sorting the unique elements: O(n × log n) in the worst case when all elements are unique
The main loop iterates n times, and for each iteration:
- bisect_left performs binary search: O(log n)
- tree.query(j) calls the BIT query operation, which loops at most O(log n) times (following the path down by removing the least significant bit)
- tree.update(j, v) calls the BIT update operation, which loops at most O(log n) times (following the path up by adding the least significant bit)
Final tree.query(len(s)): O(log n)

Overall: O(n) + O(n × log n) + O(n × log n) = O(n × log n)

Space Complexity: O(n)

The space complexity consists of:

Array arr storing transformed values: O(n)
Set and sorted array s storing unique values: O(n) in the worst case
BinaryIndexedTree with array c of size n + 1: O(n)
Additional variables and temporary storage: O(1)

Overall: O(n) + O(n) + O(n) = O(n)

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Handling of Negative Values in BIT Initialization

Pitfall: One common mistake is initializing the BIT with 0 instead of negative infinity. This causes incorrect results when all possible balanced subsequences have negative sums.

Why it happens: When querying bit.query(compressed_index), if no valid subsequence has been found yet, the BIT should return negative infinity, not 0. If initialized with 0, the algorithm might incorrectly assume there's a valid empty subsequence with sum 0.

Example scenario:

nums = [-10, -20, -30] # All negative values

With incorrect initialization (using 0), the algorithm might return 0 instead of the correct answer -10 (taking just the first element).

Solution: Always initialize the BIT with float('-inf'):

self.tree = [float('-inf')] * (n + 1) # Correct # NOT: self.tree = [0] * (n + 1) # Wrong!

2. Off-by-One Error in BIT Indexing

Pitfall: Forgetting that BIT is 1-indexed while coordinate compression returns 0-indexed positions.

Why it happens: bisect_left returns a 0-indexed position, but BIT operations expect 1-indexed positions.

Incorrect code:

compressed_index = bisect_left(sorted_unique_values, current_num - i) # Wrong! bit.update(compressed_index, new_sum) # This will fail for index 0

Solution: Always add 1 when converting from bisect result to BIT index:

compressed_index = bisect_left(sorted_unique_values, current_num - i) + 1 # Correct

3. Misunderstanding the Transformation Logic

Pitfall: Incorrectly transforming the balance condition, such as using nums[i] + i instead of nums[i] - i.

Why it happens: The balance condition nums[j] - nums[i] >= j - i needs careful algebraic manipulation:

Rearranging: nums[j] - j >= nums[i] - i
This means we need arr[i] = nums[i] - i, not nums[i] + i

Incorrect transformation:

transformed_values = [nums[i] + i for i in range(len(nums))] # Wrong!

Solution: Use the correct transformation:

transformed_values = [nums[i] - i for i in range(len(nums))] # Correct

4. Not Handling the Case of Starting a New Subsequence

Pitfall: Forgetting that a single element can form a valid balanced subsequence, leading to missing optimal solutions.

Why it happens: When calculating the maximum sum at position i, we need to consider both:

Extending a previous subsequence
Starting a new subsequence from the current element

Incorrect code:

max_previous_sum = bit.query(compressed_index) # Wrong! new_sum = max_previous_sum + current_num # Might use -inf

Solution: Take the maximum with 0 to allow starting fresh:

max_previous_sum = max(bit.query(compressed_index), 0) # Correct new_sum = max_previous_sum + current_num

5. Using Standard BIT for Sum Instead of Maximum

Pitfall: Using a standard BIT implementation that calculates prefix sums instead of prefix maximums.

Why it happens: Most BIT tutorials and implementations focus on range sum queries, but this problem requires range maximum queries.

Solution: Modify the BIT operations to use max instead of addition:

# In update method: self.tree[index] = max(self.tree[index], value) # Use max, not +=  # In query method: maximum_value = max(maximum_value, self.tree[index]) # Use max, not +=

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

Which data structure is used to implement priority queue?

Stack

Linked list

Heap

Hash table

I don't know