Skip to content

2355. Maximum Number of Books You Can Take 👍

  • Time: $O(n)$
  • Space: $O(n)$
 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
class Solution {
 public:
  long long maximumBooks(vector<int>& books) {
    // dp[i] := the maximum number of books we can take from books[0..i] with
    // taking all of books[i]
    vector<long> dp(books.size());
    stack<int> stack;  // the possible indices we can reach

    for (int i = 0; i < books.size(); ++i) {
      // We may take all of books[j], where books[j] < books[i] - (i - j).
      while (!stack.empty() &&
             books[stack.top()] >= books[i] - (i - stack.top()))
        stack.pop();
      // We can now take books[j + 1..i].
      const int j = stack.empty() ? -1 : stack.top();
      const int lastTook = books[i] - (i - j) + 1;
      if (lastTook > 1)
        // books[i] + (books[i] - 1) + ... + (books[i] - (i - j) + 1)
        dp[i] = static_cast<long>(books[i] + lastTook) * (i - j) / 2;
      else
        // 1 + 2 + ... + books[i]
        dp[i] = static_cast<long>(books[i]) * (books[i] + 1) / 2;
      if (j >= 0)
        dp[i] += dp[j];
      stack.push(i);
    }

    return ranges::max(dp);
  }
};

 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 Solution {
  public long maximumBooks(int[] books) {
    // dp[i] := the maximum number of books we can take from books[0..i] with taking all of
    // books[i]
    long[] dp = new long[books.length];
    Deque<Integer> stack = new ArrayDeque<>(); // the possible indices we can reach

    for (int i = 0; i < books.length; ++i) {
      // We may take all of books[j], where books[j] < books[i] - (i - j).
      while (!stack.isEmpty() && books[stack.peek()] >= books[i] - (i - stack.peek()))
        stack.pop();
      // We can now take books[j + 1..i].
      final int j = stack.isEmpty() ? -1 : stack.peek();
      final int lastTook = books[i] - (i - j) + 1;
      if (lastTook > 1)
        // books[i] + (books[i] - 1) + ... + (books[i] - (i - j) + 1)
        dp[i] = (long) (books[i] + lastTook) * (i - j) / 2;
      else
        // 1 + 2 + ... + books[i]
        dp[i] = (long) books[i] * (books[i] + 1) / 2;
      if (j >= 0)
        dp[i] += dp[j];
      stack.push(i);
    }

    return Arrays.stream(dp).max().getAsLong();
  }
}

 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
class Solution:
  def maximumBooks(self, books: list[int]) -> int:
    # dp[i] := the maximum the number of books we can take from books[0..i] with taking all of
    # books[i]
    dp = [0] * len(books)
    stack = []  # the possible indices we can reach

    for i, book in enumerate(books):
      # We may take all of books[j], where books[j] < books[i] - (i - j).
      while stack and books[stack[-1]] >= book - (i - stack[-1]):
        stack.pop()
      # We can now take books[j + 1..i].
      j = stack[-1] if stack else -1
      lastPicked = book - (i - j) + 1
      if lastPicked > 1:
        # book + (book - 1) + ... + (book - (i - j) + 1)
        dp[i] = (book + lastPicked) * (i - j) // 2
      else:
        # 1 + 2 + ... + book
        dp[i] = book * (book + 1) // 2
      if j >= 0:
        dp[i] += dp[j]
      stack.append(i)

    return max(dp)