Skip to content

2926. Maximum Balanced Subsequence Sum

SourceWeekly Contest 370 Q4DifficultyHardRating2448

Description

You are given a 0-indexed integer array nums.

A subsequence of nums having length k and consisting of indices i0 < i1 < ... < ik-1 is balanced if the following holds:

  • nums[ij] - nums[ij-1] >= ij - ij-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 <= 105
  • -109 <= nums[i] <= 109

Solutions

Solution 1: Dynamic Programming + Binary Indexed Tree

Thinking

Balance \(\textit{nums}[i]-\textit{nums}[j] \ge i-j\) rewrites as \(\textit{nums}[i]-i \ge \textit{nums}[j]-j\). On \(\textit{arr}[t]=\textit{nums}[t]-t\) we want a nondecreasing index sequence maximizing the sum of \(\textit{nums}\). The naive \(f[i]=\textit{nums}[i]+\max_{j<i, \textit{arr}[j]\le \textit{arr}[i]} f[j]\) (or just \(\textit{nums}[i]\)) is \(O(n^2)\) for \(n \le 10^5\).

After compressing \(\textit{arr}\), a Fenwick tree stores the best \(f\) among values up to a threshold. Scan in index order, query, then insert; a final prefix maximum is the answer.

According to the problem description, we can transform the inequality \(\textit{nums}[i] - \textit{nums}[j] \ge i - j\) into \(\textit{nums}[i] - i \ge \textit{nums}[j] - j\). Therefore, we consider defining a new array \(\textit{arr}\), where \(\textit{arr}[i] = \textit{nums}[i] - i\). A balanced subsequence satisfies that for any \(j < i\), \(\textit{arr}[j] \le \textit{arr}[i]\). The problem is transformed into selecting an increasing subsequence in \(\textit{arr}\) such that the corresponding sum in \(\textit{nums}\) is maximized.

Suppose \(i\) is the index of the last element in the subsequence, then we consider the index \(j\) of the second to last element in the subsequence. If \(\textit{arr}[j] \le \textit{arr}[i]\), we can consider whether to add \(j\) to the subsequence.

Therefore, we define \(f[i]\) as the maximum sum of \(\textit{nums}\) when the index of the last element in the subsequence is \(i\). The answer is \(\max_{i=0}^{n-1} f[i]\).

The state transition equation is:

\[ f[i] = \max(\max_{j=0}^{i-1} f[j], 0) + \textit{nums}[i] \]

where \(j\) satisfies \(\textit{arr}[j] \le \textit{arr}[i]\).

We can use a Binary Indexed Tree to maintain the maximum value of the prefix, i.e., for each \(\textit{arr}[i]\), we maintain the maximum value of \(f[i]\) in the prefix \(\textit{arr}[0..i]\).

The time complexity is \(O(n \log n)\), and the space complexity is \(O(n)\), where \(n\) is the length of the array \(\textit{nums}\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class 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))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
class 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;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class 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);
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
const inf int = 1e18

type BinaryIndexedTree struct {
    n int
    c []int
}

func NewBinaryIndexedTree(n int) BinaryIndexedTree {
    c := make([]int, n+1)
    for i := range c {
        c[i] = -inf
    }
    return BinaryIndexedTree{n: n, c: c}
}

func (bit *BinaryIndexedTree) update(x, v int) {
    for x <= bit.n {
        bit.c[x] = max(bit.c[x], v)
        x += x & -x
    }
}

func (bit *BinaryIndexedTree) query(x int) int {
    mx := -inf
    for x > 0 {
        mx = max(mx, bit.c[x])
        x -= x & -x
    }
    return mx
}

func maxBalancedSubsequenceSum(nums []int) int64 {
    n := len(nums)
    arr := make([]int, n)
    for i, x := range nums {
        arr[i] = x - i
    }
    sort.Ints(arr)
    m := 0
    for i, x := range arr {
        if i == 0 || x != arr[i-1] {
            arr[m] = x
            m++
        }
    }
    arr = arr[:m]
    tree := NewBinaryIndexedTree(m)
    for i, x := range nums {
        j := sort.SearchInts(arr, x-i) + 1
        v := max(tree.query(j), 0) + x
        tree.update(j, v)
    }
    return int64(tree.query(m))
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class BinaryIndexedTree {
    private n: number;
    private c: number[];

    constructor(n: number) {
        this.n = n;
        this.c = Array(n + 1).fill(-Infinity);
    }

    update(x: number, v: number): void {
        while (x <= this.n) {
            this.c[x] = Math.max(this.c[x], v);
            x += x & -x;
        }
    }

    query(x: number): number {
        let mx = -Infinity;
        while (x > 0) {
            mx = Math.max(mx, this.c[x]);
            x -= x & -x;
        }
        return mx;
    }
}

function maxBalancedSubsequenceSum(nums: number[]): number {
    const n = nums.length;
    const arr = Array(n).fill(0);
    for (let i = 0; i < n; ++i) {
        arr[i] = nums[i] - i;
    }
    arr.sort((a, b) => a - b);
    let m = 0;
    for (let i = 0; i < n; ++i) {
        if (i === 0 || arr[i] !== arr[i - 1]) {
            arr[m++] = arr[i];
        }
    }
    arr.length = m;
    const tree = new BinaryIndexedTree(m);
    const search = (nums: number[], x: number): number => {
        let [l, r] = [0, nums.length];
        while (l < r) {
            const mid = (l + r) >> 1;
            if (nums[mid] >= x) {
                r = mid;
            } else {
                l = mid + 1;
            }
        }
        return l;
    };
    for (let i = 0; i < n; ++i) {
        const j = search(arr, nums[i] - i) + 1;
        const v = Math.max(tree.query(j), 0) + nums[i];
        tree.update(j, v);
    }
    return tree.query(m);
}

Comments