leetcode：2926. 平衡子序列的最大和 【树状数组维护最大前缀和】

题目链接

lc2926

题目描述

题目思路

定义b[i] = nums[i] - i
目标是从b中找到一个非降子序列使得元素和最大
# b[i] = nums[i] - i
# 找到b的一个非降子序列使得元素和最大
# f[i]: 子序列最后一个数下标是i，对应的最大子序列
# f[i] = max (max f[j], 0) + nums[i] (j < i and b[j] <= b[i])
# 也就是维护f[j]的前缀最大值
# 10 ** 9: 离散化处理 + 重新标序号

ac code

# 树状数组模板（维护前缀最大值）
class BIT:
    def __init__(self, n: int):
        self.tree = [-inf] * n

    def update(self, i: int, val: int) -> None:
        while i < len(self.tree):
            self.tree[i] = max(self.tree[i], val)
            i += i & -i

    def pre_max(self, i: int) -> int:
        mx = -inf
        while i > 0:
            mx = max(mx, self.tree[i])
            i -= i & -i
        return mx


class Solution:
    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:
        n = len(nums)
        # b[i] = nums[i] - i
        # 找到b的一个非降子序列使得元素和最大
        # f[i]: 子序列最后一个数下标是i，对应的最大子序列
        # f[i] = max (max f[j], 0) + nums[i] (j < i and b[j] <= b[i])
        # 也就是维护f[j]的前缀最大值
        # 10 ** 9: 离散化处理 + 重新标序号
        b = sorted(set(x - i for i, x in enumerate(nums)))
        t = BIT(len(b) + 1) # 经典初始化
        for i, x in enumerate(nums):
            j = bisect_left(b, x - i) + 1 # 找到b[i]离散化后的位置
            f = max(t.pre_max(j), 0) + x # 计算f[i]
            t.update(j, f)
        return t.pre_max(len(b)) # 所有f的最大值
            

树状数组维护前缀最大值模版

# 树状数组模板（维护前缀最大值）
class BIT:
    def __init__(self, n: int):
        self.tree = [-inf] * n

    def update(self, i: int, val: int) -> None:
        while i < len(self.tree):
            self.tree[i] = max(self.tree[i], val)
            i += i & -i

    def pre_max(self, i: int) -> int:
        mx = -inf
        while i > 0:
            mx = max(mx, self.tree[i])
            i -= i & -i
        return mx