2926. Maximum Balanced Subsequence Sum

Problem Description

In this problem, you're given an integer array nums. The objective is to find the maximum possible sum of elements in a balanced subsequence of this array. A subsequence of the array is considered balanced if it follows these rules:

• It consists of indices (i_0, i_1, ..., i_k-1) such that i_j is less than i_{j+1} for all j.
• For each index j from 1 to k-1, the difference in values at nums[i_j] and nums[i_{j-1}] must be at least as great as the difference in their indices (i_j - i_{j-1}). Put another way, nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1}.

A subsequence with just one element is always balanced by definition.

The goal is not just to find any balanced subsequence but the one that maximizes the sum of the nums elements included in it. A subsequence in this context means a sequence that can be derived from the array by removing some (or no) elements without changing the order of the remaining elements.

Intuition

To solve this problem efficiently, we can use dynamic programming together with a Binary Indexed Tree (BIT) for optimization. Initially, the constraint nums[i_j] - nums[i_{j-1}] >= i_j - i_{j-1} might seem tricky to handle directly in terms of subsequence selection. However, we can simplify the problem with a transformation. By rearranging the inequality, we can rewrite it as nums[i_j] - i_j >= nums[i_{j-1}] - i_{j-1}.

With this insight, we can look at a new transformed array where each element is the original element subtracted by its index (arr[i] = nums[i] - i). Now, we need to find an increasing subsequence in this new array that would yield the maximum sum from the original nums array. For each element nums[i], now associated with arr[i], the subsequence can only continue with an arr[j] such that arr[j] <= arr[i] for all j < i.

We use dynamic programming where f[i] represents the maximum sum of the subsequence ending with the ith element. The result will be the maximum value of f[i] for all possible i. To efficiently calculate f[i], we must efficiently find the preceding maximum f[j] for all j where arr[j] <= arr[i].

A Binary Indexed Tree is perfect for this task because it allows us to maintain and query the maximum value of any prefix efficiently. Using a BIT, we can retrieve the maximum f[j] in logarithmic time and update the maximum value as we iterate through the array.

This transformation simplifies the original problem and allows us to utilize well-known algorithms (like a BIT) to achieve an efficient solution.

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

Solution Approach

To implement the solution, we follow a two-step approach:

First, we transform the nums array into the arr array, where for each index i, we calculate arr[i] = nums[i] - i. This transformed array serves as the basis for finding the increasing subsequence according to the restructured condition for our balanced subsequence.

Second, we use dynamic programming to find the maximum sum of such a balanced subsequence. We define f[i] as the maximum sum of the subsequence when the index i is the ending element of the subsequence. The formula to compute f[i] is based on two conditions:

1. If we include the ith element in the subsequence, then f[i] must be the maximum value of all f[j] for j < i where arr[j] <= arr[i], plus the value of nums[i].
2. If i is the first element of the subsequence, f[i] is just nums[i].

In mathematical terms, the state transition is as follows:

1f[i] = max(max(f[j] for j < i and arr[j] <= arr[i]), 0) + nums[i]

To efficiently compute this, we use a Binary Indexed Tree (BIT), also known as a Fenwick tree, which allows us to quickly perform two operations:

• Update the maximum value of f[i] at a specific transformed value arr[i].
• Query the current maximum value of f[i] for any prefix that ends with a value less than or equal to arr[i].

When iterating through each element i of the nums array, we perform a binary search on the sorted unique values of arr to find the proper index j for updating and querying in the BIT, which works since arr is sorted. This index represents a transformed space in which we're updating and querying the maximum sums.

As a result, the BIT maintains the maximum sums for all processed indices in a compressed form, and we are able to compute f[i] efficiently for every element. After iterating through all elements, the maximum value stored in BIT will represent the maximum possible sum of all balanced subsequences in the original nums array.

The final return value is obtained by querying the BIT for the global maximum, which represents the maximum sum of elements in a balanced subsequence, and this sum is the answer we need.

By addressing both the need for quick updates and maximum queries, the combination of dynamic programming with the Binary Indexed Tree provides an optimized solution to the problem.

Example Walkthrough

Let's say our input array nums = [3, 1, 2, 5, 3]. We will illustrate the solution approach using this array.

1. Transform the input array into the arr array by calculating arr[i] = nums[i] - i for each element. Our transformed arr array looks like this:

   1nums:  [3,  1,  2,  5,  3]
   2i:      0   1   2   3   4
   3arr:   [3,  0,  0,  2, -1]

2. Create an array f to hold the maximum sums of subsequences ending at each index i. Initially, f will just be a copy of nums since each element can form a subsequence on its own:

   1f: [3,  1,  2,  5,  3]

3. Next, we sort the unique values of arr and use this sorted array to determine the indices for updates and queries in the Binary Indexed Tree (BIT). Sorted arr will be [0, 0, -1, 2, 3].

4. Iterate through each element i of the nums array:

   • For i = 0, arr[0] = 3. Since it's the first element, we just initialize the BIT at arr[0] with f[0] = nums[0] = 3.
   • For i = 1, arr[1] = 0. arr[1] is less than arr[0], so we don't update f[1]. The f[1] remains 1.
   • For i = 2, arr[2] = 0. arr[2] is equal to arr[1] and also less than arr[0], we find the maximum between f[2] and f[1]. f[2] = max(f[2], f[1]) = max(2, 1) = 2.
   • For i = 3, arr[3] = 2. All previous entries in arr are less than or equal to arr[3], thus we have to choose the maximum sum f of those. This is obtained from the BIT, let's assume it returns f[j] for the best j < 3. The updated f[3] = max(f[j]) + nums[3] = max(3) + 5 = 8.
   • For i = 4, arr[4] = -1. arr[4] is less than all previous entries except for itself, so we can't extend any subsequence that ends before it. f[4] remains 3.

5. Now the f array is updated and we keep track of the maximum value encountered at each step.

6. The final answer is the maximum value in the f array which represents the maximum possible sum of elements in a balanced subsequence. In this case, max(f) = 8, since f[3] = 8 is the maximum value we encountered corresponding to the subsequence of nums that produced it [3, 5].

Solution Implementation

Time and Space Complexity

The provided code implements a Binary Indexed Tree (BIT), also known as a Fenwick Tree, which is used to perform operations on prefix sums efficiently. The time complexity and space complexity of the key operations in the code are as follows:

Time Complexity

The time complexity of the code is dominated by two main operations: update and query on the Binary Indexed Tree, and sorting the distinct elements of the array. The BIT operations update() and query() each have a time complexity of O(log n) because each iteration goes up or down the tree, which has a height proportional to the logarithm of the number of elements (n).

1. Sorting distinct elements: To sort the s array, which contains the distinct values of arr, it takes O(k log k) time, where k is the number of distinct values in arr. Since k <= n, the sorting step can also be bounded by O(n log n).

2. Loop with BIT operations: The for loop iterates n times and performs the update or query operation once per iteration. As each BIT operation is O(log n), the total time complexity for the loop is O(n log n).

Therefore, the combined time complexity of the algorithm is O(n log n) + O(n log n), which simplifies to O(n log n).

Space Complexity

The space complexity is driven by the space allocated for the BIT and the sorted set of elements. The BIT needs an array c of size n + 1, where n is the number of elements in the given input nums. The sorted array s stores distinct elements of the transformed array arr, which, in the worst case, can be n elements if all arr values are unique.

Hence, the space complexity of the algorithm is O(n) for storing the BIT and O(n) for storing the distinct elements, which results in a total space complexity of O(2n). However, this is simplified to O(n) because constant factors are dropped in big-O notation.

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