Maximum Balanced Subsequence Sum

Hard

#2606Time: O(n^2), where `n` is the length of the `nums` array. The nested loops dominate the runtime.Space: O(n) to store the `dp` array.

Patterns

Dynamic Programming

Algorithms

Binary Search

Data structures

Array Binary Indexed Tree Segment Tree

Prompt

You are given a 0-indexed integer array nums.

A subsequence of nums having length k and consisting of indices i₀ < i₁ < ... < i_k-1 is balanced if the following holds:

nums[i_j] - nums[i_j-1] >= i_j - i_j-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.

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⁵
-10⁹ <= nums[i] <= 10⁹

Approaches

2 approaches with complexity analysis and trade-offs.

This approach uses dynamic programming to solve the problem. First, we transform the problem's condition. The condition nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1} can be rewritten as nums[i_j] - i_j >= nums[i_{j-1}] - i_{j-1}. Let's define a new array b where b[i] = nums[i] - i. The problem now is to find a subsequence with non-decreasing b values that has the maximum sum of corresponding nums values.

We define dp[i] as the maximum sum of a balanced subsequence ending at index i. To compute dp[i], we look at all previous indices j < i. If b[j] <= b[i], it means we can extend a balanced subsequence ending at j with the element at i. We choose the j that gives the maximum dp[j] to maximize the new sum.

Algorithm

Create a dp array of size n, where n is the length of nums.
Initialize a variable max_sum to the smallest possible long value.
Loop i from 0 to n-1:
- Initialize maxPrevSum = 0.
- Loop j from 0 to i-1:
  - If (long)nums[j] - j <= (long)nums[i] - i, update maxPrevSum = Math.max(maxPrevSum, dp[j]).
- Set dp[i] = nums[i] + maxPrevSum.
- Update max_sum = Math.max(max_sum, dp[i]).
Return max_sum.

Walkthrough

The recurrence relation is dp[i] = nums[i] + max({0} U {dp[j] | 0 <= j < i and b[j] <= b[i]}), where b[k] = nums[k] - k. We iterate through each element of the array from left to right. For each element nums[i], we calculate the maximum possible sum of a balanced subsequence ending at this index. This is done by finding the maximum sum of a valid preceding subsequence and adding nums[i] to it. A preceding subsequence ending at index j is valid if j < i and nums[j] - j <= nums[i] - i. If no such valid preceding subsequence exists or if all of them have a sum less than or equal to zero, we can start a new subsequence with just nums[i], which means we add 0 to nums[i].

The algorithm is as follows:

Create an array dp of the same size as nums, where dp[i] will store the maximum sum of a balanced subsequence ending at index i.
Iterate through the nums array with index i from 0 to n-1.
For each i, initialize a variable maxPrevSum to 0.
Start an inner loop with index j from 0 to i-1.
Inside the inner loop, check if the balanced condition is met: (long)nums[j] - j <= (long)nums[i] - i.
If the condition is met, update maxPrevSum = Math.max(maxPrevSum, dp[j]).
After the inner loop, calculate dp[i] = nums[i] + maxPrevSum.
The final answer is the maximum value found in the dp array.

1import java.util.Arrays;2 3class Solution {4    public long maxBalancedSubsequenceSum(int[] nums) {5        int n = nums.length;6        long[] dp = new long[n];7        long maxSum = Long.MIN_VALUE;8 9        for (int i = 0; i < n; i++) {10            long maxPrevSum = 0;11            for (int j = 0; j < i; j++) {12                if ((long)nums[j] - j <= (long)nums[i] - i) {13                    maxPrevSum = Math.max(maxPrevSum, dp[j]);14                }15            }16            dp[i] = (long)nums[i] + maxPrevSum;17            maxSum = Math.max(maxSum, dp[i]);18        }19 20        return maxSum;21    }22}

Complexity

Time

O(n^2), where `n` is the length of the `nums` array. The nested loops dominate the runtime.

Space

O(n) to store the `dp` array.

Trade-offs

Pros

Simple to understand and implement.
Correctly solves the problem based on the DP formulation.

Cons

Inefficient due to the nested loops.
Time complexity of O(n^2) is too slow for the given constraints (n <= 10^5), leading to a "Time Limit Exceeded" error on large inputs.

The O(n^2) DP approach can be optimized. The bottleneck is the inner loop which finds the maximum dp[j] among all j < i satisfying b[j] <= b[i] (where b[i] = nums[i] - i). This subproblem is a range maximum query. We can use a data structure to perform this query efficiently. A Fenwick Tree (also known as a Binary Indexed Tree or BIT) is suitable for this.

Since the values of b[i] can be large and negative, we first need to perform coordinate compression on them. We map each unique b[i] value to a smaller, non-negative integer rank. Then, we use a Fenwick Tree that supports range maximum queries. As we iterate through i from 0 to n-1, we query the Fenwick Tree for the maximum dp value for all b values less than or equal to the current b[i], calculate the new dp[i], and then update the Fenwick Tree with this new value at the rank corresponding to b[i].

Algorithm

Create an array b where b[i] = nums[i] - i.
Perform coordinate compression on b:
- Get the unique values of b.
- Sort these unique values.
- Create a map from each unique value to its sorted index (rank).
Initialize a Fenwick Tree ft of size equal to the number of unique b values. The tree should support range maximum queries and its neutral element should be 0.
Initialize maxSum to Long.MIN_VALUE.
Loop i from 0 to n-1:
- Get the rank of b[i] from the map.
- Query ft up to this rank to get prevMaxSum. This value represents max(0, max_{j<i, b[j]<=b[i]} dp[j]).
- Calculate currentSum = nums[i] + prevMaxSum.
- Update ft at the rank of b[i] with currentSum.
- Update maxSum = max(maxSum, currentSum).
Return maxSum.

Walkthrough

The core idea is to maintain a data structure that, for each possible value of v = nums[j] - j, stores the maximum balanced subsequence sum ending with that v. When we process nums[i], we need to find the maximum sum among all subsequences that can precede it. These are the ones ending at j < i where nums[j] - j <= nums[i] - i.

The algorithm is as follows:

Create a new array b where b[i] = nums[i] - i.
The values in b can be large, so we perform coordinate compression. We collect all unique values from b, sort them, and create a map from each value to its rank.
Initialize a Fenwick Tree (BIT) designed for range maximum queries. The size of the BIT will be the number of unique values in b. Its neutral element for queries should be 0, representing the choice to start a new subsequence.
Iterate through the nums array from i = 0 to n-1.
For each i, find the rank of b[i].
Query the BIT to find the maximum value in the range up to b[i]'s rank. This gives max(0, max_{j<i, b[j]<=b[i]} dp[j]). Let's call this prevMaxSum.
Calculate the maximum sum for a balanced subsequence ending at i: currentSum = nums[i] + prevMaxSum.
Update the BIT at b[i]'s rank with currentSum.
Keep track of the overall maximum sum found so far.
The maximum of all currentSum values calculated is the result.

1import java.util.*;2 3class Solution {4    static class FenwickTree {5        private long[] tree;6        private int size;7        private static final long NEUTRAL_ELEMENT = 0L;8 9        public FenwickTree(int size) {10            this.size = size;11            this.tree = new long[size + 1];12        }13 14        public void update(int idx, long val) {15            idx++; // 1-based index16            while (idx <= size) {17                tree[idx] = Math.max(tree[idx], val);18                idx += idx & -idx;19            }20        }21 22        public long query(int idx) {23            idx++; // 1-based index24            long maxVal = NEUTRAL_ELEMENT;25            while (idx > 0) {26                maxVal = Math.max(maxVal, tree[idx]);27                idx -= idx & -idx;28            }29            return maxVal;30        }31    }32 33    public long maxBalancedSubsequenceSum(int[] nums) {34        int n = nums.length;35        int[] b = new int[n];36        for (int i = 0; i < n; i++) {37            b[i] = nums[i] - i;38        }39 40        Set<Integer> uniqueBSet = new HashSet<>();41        for (int val : b) {42            uniqueBSet.add(val);43        }44        List<Integer> sortedUniqueB = new ArrayList<>(uniqueBSet);45        Collections.sort(sortedUniqueB);46        47        Map<Integer, Integer> rankMap = new HashMap<>();48        for (int i = 0; i < sortedUniqueB.size(); i++) {49            rankMap.put(sortedUniqueB.get(i), i);50        }51 52        FenwickTree ft = new FenwickTree(sortedUniqueB.size());53        long maxSum = Long.MIN_VALUE;54 55        for (int i = 0; i < n; i++) {56            int currentRank = rankMap.get(b[i]);57            long prevMaxSum = ft.query(currentRank);58            59            long currentSum = (long)nums[i] + prevMaxSum;60            61            ft.update(currentRank, currentSum);62            maxSum = Math.max(maxSum, currentSum);63        }64 65        return maxSum;66    }67}

Complexity

Time

O(n log n). Coordinate compression takes O(n log n). The main loop runs `n` times, with each Fenwick Tree operation taking O(log n) time.

Space

O(n). This includes space for the `b` array, coordinate compression map, and the Fenwick Tree, each of which can be up to O(n) in size.

Trade-offs

Pros

Highly efficient with O(n log n) time complexity, which passes the constraints.
Solves a more general class of problems (Maximum Sum Non-decreasing Subsequence on a transformed array).

Cons

More complex to implement than the brute-force DP.
Requires understanding of advanced data structures like Fenwick Trees and coordinate compression.

Solutions

Solution

1class BinaryIndexedTree { private int n ; private long [] c ; private final long inf = 1L << 60 ; public BinaryIndexedTree ( int n ) { this . n = n ; c = new long [ n + 1 ]; Arrays . fill ( c , - inf ); } public void update ( int x , long v ) { while ( x <= n ) { c [ x ] = Math . max ( c [ x ], v ); x += x & - x ; } } public long query ( int x ) { long mx = - inf ; while ( x > 0 ) { mx = Math . max ( mx , c [ x ]); x -= x & - x ; } return mx ; } } class Solution { public long maxBalancedSubsequenceSum ( int [] nums ) { int n = nums . length ; int [] arr = new int [ n ]; for ( int i = 0 ; i < n ; ++ i ) { arr [ i ] = nums [ i ] - i ; } Arrays . sort ( arr ); int m = 0 ; for ( int i = 0 ; i < n ; ++ i ) { if ( i == 0 || arr [ i ] != arr [ i - 1 ]) { arr [ m ++] = arr [ i ]; } } BinaryIndexedTree tree = new BinaryIndexedTree ( m ); for ( int i = 0 ; i < n ; ++ i ) { int j = search ( arr , nums [ i ] - i , m ) + 1 ; long v = Math . max ( tree . query ( j ), 0 ) + nums [ i ]; tree . update ( j , v ); } return tree . query ( m ); } private int search ( int [] nums , int x , int r ) { int l = 0 ; while ( l < r ) { int mid = ( l + r ) >> 1 ; if ( nums [ mid ] >= x ) { r = mid ; } else { l = mid + 1 ; } } return l ; } }

1class BinaryIndexedTree { private: int n ; vector < long long > c ; const long long inf = 1e18 ; public: BinaryIndexedTree ( int n ) { this -> n = n ; c . resize ( n + 1 , - inf ); } void update ( int x , long long v ) { while ( x <= n ) { c [ x ] = max ( c [ x ], v ); x += x & - x ; } } long long query ( int x ) { long long mx = - inf ; while ( x > 0 ) { mx = max ( mx , c [ x ]); x -= x & - x ; } return mx ; } }; class Solution { public: long long maxBalancedSubsequenceSum ( vector < int >& nums ) { int n = nums . size (); vector < int > arr ( n ); for ( int i = 0 ; i < n ; ++ i ) { arr [ i ] = nums [ i ] - i ; } sort ( arr . begin (), arr . end ()); arr . erase ( unique ( arr . begin (), arr . end ()), arr . end ()); int m = arr . size (); BinaryIndexedTree tree ( m ); for ( int i = 0 ; i < n ; ++ i ) { int j = lower_bound ( arr . begin (), arr . end (), nums [ i ] - i ) - arr . begin () + 1 ; long long v = max ( tree . query ( j ), 0LL ) + nums [ i ]; tree . update ( j , v ); } return tree . query ( m ); } };

1class BinaryIndexedTree : def __init__ ( self , n : int ): self . n = n self . c = [ - inf ] * ( n + 1 ) def update ( self , x : int , v : int ): while x <= self . n : self . c [ x ] = max ( self . c [ x ], v ) x += x & - x def query ( self , x : int ) -> int : mx = - inf while x : mx = max ( mx , self . c [ x ]) x -= x & - x return mx class Solution : def maxBalancedSubsequenceSum ( self , nums : List [ int ]) -> int : arr = [ x - i for i , x in enumerate ( nums )] s = sorted ( set ( arr )) tree = BinaryIndexedTree ( len ( s )) for i , x in enumerate ( nums ): j = bisect_left ( s , x - i ) + 1 v = max ( tree . query ( j ), 0 ) + x tree . update ( j , v ) return tree . query ( len ( s ))

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem