Maximum balanced subsequence sum

Array Binary Search Dynamic Programming Binary Indexed Tree Segment Tree

25.8% acceptance

Feb 25, 2026

289

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: 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.

Solution

Rust

Time O(n log n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn max_balanced_subsequence_sum(nums: Vec<i32>) -> i64 {
    // Condition: nums[j] - nums[i] >= j - i  ↔  nums[j]-j >= nums[i]-i
    // Let a[i] = nums[i] - i. Need non-decreasing subsequence of a[].
    // Maximize sum of nums[i] over the selected indices.
    // dp[i] = max sum of balanced subseq ending at i
    // dp[i] = nums[i] + max(0, max{dp[j] : j < i, a[j] <= a[i]})

    let n = nums.len();
    let a: Vec<i64> = (0..n).map(|i| nums[i] as i64 - i as i64).collect();

    // Coordinate compress a[]
    let mut sorted_a = a.clone();
    sorted_a.sort_unstable();
    sorted_a.dedup();
    let rank = |x: i64| -> usize {
      sorted_a.partition_point(|&v| v < x)
    };

    let m = sorted_a.len();
    // BIT for prefix max
    let mut bit = vec![i64::MIN / 2; m + 1];

    fn bit_update(bit: &mut Vec<i64>, mut i: usize, v: i64) {
      i += 1;
      while i < bit.len() {
        if v > bit[i] { bit[i] = v; }
        i += i & i.wrapping_neg();
      }
    }
    fn bit_query(bit: &Vec<i64>, mut i: usize) -> i64 {
      i += 1;
      let mut res = i64::MIN / 2;
      while i > 0 {
        if bit[i] > res { res = bit[i]; }
        i -= i & i.wrapping_neg();
      }
      res
    }

    let mut ans = i64::MIN;
    for i in 0..n {
      let r = rank(a[i]);
      let prev_best = bit_query(&bit, r);
      let dp_i = nums[i] as i64 + if prev_best > 0 { prev_best } else { 0 };
      let dp_i = dp_i.max(nums[i] as i64); // always at least take nums[i] alone
      ans = ans.max(dp_i);
      bit_update(&mut bit, r, dp_i);
    }
    ans
  }
}