Problem statement

https://leetcode.com/problems/4-keys-keyboard/

Solution

Let dp[i] be the maximum number of A we can have, i operations. Then the last operation can be key 1: then we have dp[i-1] + 1 letters A; or we can apply key 2, key 3 + several times key 4, then we can have dp[i-k]*(k-1) letters A, because we can increase number of letters in k-1 times, using k operations. Complexity will be O(n^2). It can be optimized and shown that k is always equal to 4 or 5, so complexity can be reduced to O(n) as in the following code.

Complexity

Time complexity is O(n), space complexity is O(n), which can be reduced to O(1).

Code

class Solution:
    def maxA(self, N):
        dp = [0] * (N+1)
        dp[1] = 1
        for i in range(2, N+1):
            dp[i] = max(dp[i-1] + 1, dp[i])
            for k in range(4, min(6,i-1)):
                dp[i] = max(dp[i], dp[i-k]*(k-1))
        
        return dp[-1]