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:

f[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.

Ready to land your dream job?

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

Start Evaluator

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:

    nums:  [3,  1,  2,  5,  3]
    i:      0   1   2   3   4
    arr:   [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:

    f: [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

1from bisect import bisect_left
2from typing import List
3
4class FenwickTree:
5    # A class representing a Fenwick Tree (Binary Indexed Tree)
6    def __init__(self, size: int):
7        self.size = size
8        self.tree = [-float('inf')] * (size + 1)  # Initialize the tree with negative infinity
9
10    def update(self, index: int, value: int):
11        # Update the tree with a max value at a given index
12        while index <= self.size:
13            self.tree[index] = max(self.tree[index], value)
14            index += index & -index  # Move to the next index to be updated
15
16    def query(self, index: int) -> int:
17        # Query the maximum value up to a given index
18        max_value = -float('inf')
19        while index > 0:
20            max_value = max(max_value, self.tree[index])
21            index -= index & -index  # Move to the previous index to be queried
22        return max_value
23
24
25class Solution:
26    def max_balanced_subsequence_sum(self, nums: List[int]) -> int:
27        # A function to compute the maximum balanced subsequence sum from a list of numbers
28        adjusted_values = [x - i for i, x in enumerate(nums)]
29        sorted_unique_values = sorted(set(adjusted_values))
30        tree = FenwickTree(len(sorted_unique_values))
31
32        for i, num in enumerate(nums):
33            index = bisect_left(sorted_unique_values, num - i) + 1
34            cum_max_value = max(tree.query(index), 0) + num  # Compare the maximum with zero to avoid negative results
35            tree.update(index, cum_max_value)  # Update the tree with the new cumulative maximum
36
37        return tree.query(len(sorted_unique_values))  # Query the maximum value of the entire tree
38
1import java.util.Arrays;
2
3// BinaryIndexedTree for range query and point updates.
4class BinaryIndexedTree {
5    private int size; // Size of the array for which Binary Indexed Tree is built.
6    private long[] tree; // Internal data structure to store the tree.
7    private final long NEG_INF = Long.MIN_VALUE; // Representation for negative infinity.
8
9    // Constructor to initialize the tree.
10    public BinaryIndexedTree(int size) {
11        this.size = size;
12        tree = new long[size + 1]; // Initialized with size+1 because the index starts from 1.
13        Arrays.fill(tree, NEG_INF); // Initially, fill the tree with negative infinity.
14    }
15
16    // Update function to apply a given value to the tree.
17    public void update(int index, long value) {
18        while (index <= size) {
19            tree[index] = Math.max(tree[index], value); // Take the maximum value at this index.
20            index += index & -index; // Move to the next index to update.
21        }
22    }
23
24    // Query function to get the maximum value up to a specific index.
25    public long query(int index) {
26        long max = NEG_INF; // Start with negative infinity as the maximum.
27        while (index > 0) {
28            max = Math.max(max, tree[index]); // Update max if a greater value is found.
29            index -= index & -index; // Move index to parent.
30        }
31        return max; // Return the maximum value found.
32    }
33}
34
35class Solution {
36    // Method to find the maximum balanced subsequence sum in nums.
37    public long maxBalancedSubsequenceSum(int[] nums) {
38        int n = nums.length;
39        int[] balancedArray = new int[n]; // Array to hold the balanced values.
40        // Calculate balanced value for each element (value - its index).
41        for (int i = 0; i < n; ++i) {
42            balancedArray[i] = nums[i] - i;
43        }
44        Arrays.sort(balancedArray); // Sort the balanced values.
45      
46        // Perform compression on the balanced values. Unique values only.
47        int uniqueCount = 0;
48        for (int i = 0; i < n; ++i) {
49            if (i == 0 || balancedArray[i] != balancedArray[i - 1]) {
50                balancedArray[uniqueCount++] = balancedArray[i];
51            }
52        }
53
54        // Initialize Binary Indexed Tree with the number of unique balanced values.
55        BinaryIndexedTree tree = new BinaryIndexedTree(uniqueCount);
56        for (int i = 0; i < n; ++i) {
57            // Find the index in compressed array for each number.
58            int compressedIndex = search(balancedArray, nums[i] - i, uniqueCount) + 1;
59            // Get the best previous sum and add current value.
60            long updatedSum = Math.max(tree.query(compressedIndex), 0) + nums[i];
61            // Update tree with the new sum.
62            tree.update(compressedIndex, updatedSum);
63        }
64        // Query for the maximum balanced subsequence sum.
65        return tree.query(uniqueCount);
66    }
67
68    // Binary search to find the index of x in nums array within range r.
69    private int search(int[] nums, int x, int r) {
70        int l = 0; // Left boundary of the search.
71        while (l < r) {
72            int mid = (l + r) >> 1; // Middle index.
73            if (nums[mid] >= x) {
74                r = mid; // Search in the left half if mid value is greater or equal to x.
75            } else {
76                l = mid + 1; // Search in the right half otherwise.
77            }
78        }
79        return l; // Return the position where x fits or should be.
80    }
81}
82
1#include <vector>
2#include <algorithm>
3using namespace std;
4
5class BinaryIndexedTree {
6private:
7    int size_; // Size of the array
8    vector<long long> tree_; // The underlying data structure
9    const long long INF = 1e18; // An arbitrary large value to represent negative infinity
10  
11public:
12    // Constructor to initialize the binary indexed tree with given size
13    BinaryIndexedTree(int size) : size_(size), tree_(size + 1, -INF) { }
14
15    // Updates the tree by setting the value at 'index' to the maximum of its current value and 'value'
16    void update(int index, long long value) {
17        while (index <= size_) {
18            tree_[index] = max(tree_[index], value);
19            index += index & -index; // Traverse to parent node
20        }
21    }
22
23    // Queries the maximum value in the range 1...index
24    long long query(int index) {
25        long long max_value = -INF;
26        while (index > 0) {
27            max_value = max(max_value, tree_[index]);
28            index -= index & -index; // Traverse to child node
29        }
30        return max_value;
31    }
32};
33
34class Solution {
35public:
36    // Function to compute the maximum balanced subsequence sum
37    long long maxBalancedSubsequenceSum(vector<int>& nums) {
38        int n = nums.size();
39        vector<int> processed_nums(n);
40      
41        // Preprocess the numbers by offsetting with their indices
42        for (int i = 0; i < n; ++i) {
43            processed_nums[i] = nums[i] - i;
44        }
45      
46        // Sort and deduplicate the processed numbers
47        sort(processed_nums.begin(), processed_nums.end());
48        processed_nums.erase(unique(processed_nums.begin(), processed_nums.end()), processed_nums.end());
49      
50        int m = processed_nums.size();
51        BinaryIndexedTree tree(m);
52      
53        // Process original numbers to update the tree with the maximum value at each point
54        for (int i = 0; i < n; ++i) {
55            int idx = lower_bound(processed_nums.begin(), processed_nums.end(), nums[i] - i) - processed_nums.begin() + 1;
56            long long value = max(tree.query(idx), 0LL) + nums[i];
57            tree.update(idx, value);
58        }
59      
60        // Query the maximum value in the tree, which represents the maximum balanced subsequence sum
61        return tree.query(m);
62    }
63};
64
1const N_MAX = 100010;  // Adjust this value as necessary for your problem constraints
2let treeSize = 0;
3const tree = Array(N_MAX).fill(-Infinity);
4
5// Updates the BIT with the value 'v' at position 'index'
6function update(index: number, v: number): void {
7    while (index <= treeSize) {
8        tree[index] = Math.max(tree[index], v);
9        index += index & -index;
10    }
11}
12
13// Queries the BIT to find the maximum value up to position 'index'
14function query(index: number): number {
15    let mx = -Infinity;
16    while (index > 0) {
17        mx = Math.max(mx, tree[index]);
18        index -= index & -index;
19    }
20    return mx;
21}
22
23// Performs binary search on the array 'nums' to find the position of value 'x'
24function search(nums: number[], x: number): number {
25    let l = 0, r = nums.length;
26    while (l < r) {
27        const mid = Math.floor((l + r) / 2);
28        if (nums[mid] >= x) {
29            r = mid;
30        } else {
31            l = mid + 1;
32        }
33    }
34    return l;
35}
36
37// Computes the maximum balanced subsequence sum of the given array 'nums'
38function maxBalancedSubsequenceSum(nums: number[]): number {
39    const n = nums.length;
40    const modifiedNums = Array(n).fill(0);
41    for (let i = 0; i < n; ++i) {
42        modifiedNums[i] = nums[i] - i;
43    }
44
45    // Sort the modified nums array
46    modifiedNums.sort((a, b) => a - b);
47
48    let m = 0; // Unique count of modified values
49    // Deduplicate the sorted array
50    for (let i = 0; i < n; ++i) {
51        if (i === 0 || modifiedNums[i] !== modifiedNums[i - 1]) {
52            modifiedNums[m++] = modifiedNums[i];
53        }
54    }
55    modifiedNums.length = m; // Truncate the array to unique values count
56
57    treeSize = m; // Set the size of the BIT
58    for (let i = 0; i < n; ++i) {
59        const j = search(modifiedNums, nums[i] - i) + 1;
60        const val = Math.max(query(j), 0) + nums[i];
61        update(j, val);
62    }
63    return query(m);
64}
65

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.


Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

How many ways can you arrange the three letters A, B and C?


Recommended Readings

Want a Structured Path to Master System Design Too? Don’t Miss This!