Python recursive with memoization

  • 0

    On any day you can be in one of these 3 states:

    • flat - you don't have a position
    • long - you have a position
    • cool - you don't have a position and you are cooling down

    With this in mind, here is the solution:

    def maxProfit(self, prices):
        mem = {}
        def maxp(i, s):
            """Return the max profit you can get starting on day i with state s."""
            if i >= len(prices): return 0
            if (i, s) not in mem:
                if s == 'flat':
                    mem[(i, s)] = max(
                        maxp(i+1, 'long'),          # buy
                        maxp(i+1, 'flat'))          # don't buy
                elif s == 'long':
                    p = prices[i] - prices[i-1]
                    mem[(i, s)] = p + max(
                        maxp(i+1, 'cool'),          # sell
                        maxp(i+1, 'long'))          # don't sell
                else: # 'cool'
                    mem[(i, s)] = maxp(i+1, 'flat') # don't buy
            return mem[(i, s)]
        return maxp(i=0, s='flat')

    Note that you can replace maxp(i+1, 'cool') with maxp(i+2, 'flat') and you will not need the cool state any more :)

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.