2926. Maximum Balanced Subsequence Sum

• Time: $O(n\log n)$
• Space: $O(n)$

C++JavaPython

class FenwickTree {
 public:
  FenwickTree(int n) : vals(n + 1) {}

  // Updates the maximum sum of subsequence ending in (i - 1) with `val`.
  void maximize(int i, long val) {
    while (i < vals.size()) {
      vals[i] = max(vals[i], val);
      i += lowbit(i);
    }
  }

  // Returns the maximum sum of subsequence ending in (i - 1).
  long get(int i) const {
    long res = 0;
    while (i > 0) {
      res = max(res, vals[i]);
      i -= lowbit(i);
    }
    return res;
  }

 private:
  vector<long> vals;

  static int lowbit(int i) {
    return i & -i;
  }
};

class Solution {
 public:
  long long maxBalancedSubsequenceSum(vector<int>& nums) {
    // Let's define maxSum[i] := subsequence with the maximum sum ending in i
    // By observation:
    //     nums[i] - nums[j] >= i - j
    //  => nums[i] - i >= nums[j] - j
    //  So, if nums[i] - i >= nums[j] - j, where i > j,
    //  maxSum[i] = max(maxSum[i], maxSum[j] + nums[i])
    long ans = LONG_MIN;
    FenwickTree tree(nums.size());

    for (const auto& [_, i] : getPairs(nums)) {
      const long subseqSum = tree.get(i) + nums[i];
      tree.maximize(i + 1, subseqSum);
      ans = max(ans, subseqSum);
    }

    return ans;
  }

 private:
  vector<pair<int, int>> getPairs(const vector<int>& nums) {
    vector<pair<int, int>> pairs;
    for (int i = 0; i < nums.size(); ++i)
      pairs.emplace_back(nums[i] - i, i);
    ranges::sort(pairs);
    return pairs;
  }
};

class FenwickTree {
  public FenwickTree(int n) {
    vals = new long[n + 1];
  }

  // Updates the maximum sum of subsequence ending in (i - 1) with `val`.
  public void maximize(int i, long val) {
    while (i < vals.length) {
      vals[i] = Math.max(vals[i], val);
      i += lowbit(i);
    }
  }

  // Returns the maximum sum of subsequence ending in (i - 1).
  public long get(int i) {
    long res = 0;
    while (i > 0) {
      res = Math.max(res, vals[i]);
      i -= lowbit(i);
    }
    return res;
  }

  private long[] vals;

  private static int lowbit(int i) {
    return i & -i;
  }
}

class Solution {
  public long maxBalancedSubsequenceSum(int[] nums) {
    // Let's define maxSum[i] := subsequence with the maximum sum ending in i
    // By observation:
    //     nums[i] - nums[j] >= i - j
    //  => nums[i] - i >= nums[j] - j
    //  So, if nums[i] - i >= nums[j] - j, where i > j,
    //  maxSum[i] = max(maxSum[i], maxSum[j] + nums[i])
    long ans = Long.MIN_VALUE;
    FenwickTree tree = new FenwickTree(nums.length);

    for (Pair<Integer, Integer> pair : getPairs(nums)) {
      final int i = pair.getValue();
      final long subseqSum = tree.get(i) + nums[i];
      tree.maximize(i + 1, subseqSum);
      ans = Math.max(ans, subseqSum);
    }

    return ans;
  }

  private List<Pair<Integer, Integer>> getPairs(int[] nums) {
    List<Pair<Integer, Integer>> pairs = new ArrayList<>();
    for (int i = 0; i < nums.length; ++i)
      pairs.add(new Pair<>(nums[i] - i, i));
    pairs.sort((p1, p2) -> p1.getKey() - p2.getKey());
    return pairs;
  }
}

class FenwickTree:
  def __init__(self, n: int):
    self.vals = [0] * (n + 1)

  def maximize(self, i: int, val: int) -> None:
    """Updates the maximum sum of subsequence ending in (i - 1) with `val`."""
    while i < len(self.vals):
      self.vals[i] = max(self.vals[i], val)
      i += FenwickTree.lowbit(i)

  def get(self, i: int) -> int:
    """Returns the maximum sum of subsequence ending in (i - 1)."""
    res = 0
    while i > 0:
      res = max(res, self.vals[i])
      i -= FenwickTree.lowbit(i)
    return res

  @staticmethod
  def lowbit(i: int) -> int:
    return i & -i


class Solution:
  def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:
    # Let's define maxSum[i] := subsequence with the maximum sum ending in i
    # By observation:
    #    nums[i] - nums[j] >= i - j
    # => nums[i] - i >= nums[j] - j
    # So, if nums[i] - i >= nums[j] - j, where i > j,
    # maxSum[i] = max(maxSum[i], maxSum[j] + nums[i])
    ans = -math.inf
    tree = FenwickTree(len(nums))

    for _, i in sorted([(num - i, i) for i, num in enumerate(nums)]):
      subseqSum = tree.get(i) + nums[i]
      tree.maximize(i + 1, subseqSum)
      ans = max(ans, subseqSum)

    return ans