Leetcode: 2355. Maximum Number of Books You Can Take

from typing import List


class Solution:
    def maximumBooks(self, books: List[int]) -> int:
        st = []
        best = [0] * len(books)
        pa = lambda n: (n * (n + 1)) // 2
        ans = 0
        for i, b in enumerate(books):
            while st and books[st[-1]] >= b - (i - st[-1]):
                st.pop()
            if not st:
                best[i] = pa(b) if b <= i else pa(b) - pa(b - i - 1)
            else:
                best[i] = pa(b) if b <= i - st[-1] else pa(b) - pa(b - (i - st[-1]))
                best[i] += best[st[-1]]
            if best[i] > ans:
                ans = best[i]
                tmp = i
            st.append(i)
        return ans