Leetcode 1997 First Day Where You Have Been in All the Rooms

Problem statement

https://leetcode.com/problems/first-day-where-you-have-been-in-all-the-rooms/

Solution

The idea is that we need to go back a lot of times. Define by dp[i] is number of days to reach cell i. We can only reach it from the cell i-1, so we need at least dp[i-1] steps. Then it will lead us to dp[A[i-1]] cell, because we visited i-1 for the first time. Then we need to reach i-1 from A[i-1] again. And finally we need one more step.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def firstDayBeenInAllRooms(self, A):
        n, M = len(A), 10**9 + 7
        dp = [0]*n
        for i in range(1, n):
            dp[i] = (2*dp[i-1] - dp[A[i-1]] + 2) % M

        return dp[-1]