Problem statement

https://leetcode.com/problems/maximum-vacation-days/

Solution

Let us introduce dp(week, city) state, where week is current week number and city is current city we are in. Then we just need to check all cities from which we can reach city and return maximum. If we have zero week, we return days[city][0] if we can reach city from starting city 0 or -inf in opposite case.

Complexity

Time complexity is O(n^2k), because we have nk states and for each we need to do at most n iterations. Space complexity is O(nk). We also have G, which has inverse edges so we can traverse them fast.

Code

class Solution:
    def maxVacationDays(self, flights, days):
        G = defaultdict(set)
        n, k = len(flights), len(days[0])-1
        for i, j in product(range(n), range(n)):
            if flights[i][j] == 1 or i==j: G[j].add(i)

        @lru_cache(None)
        def dp(week, city):
            if week == 0:  
                return days[city][0] if 0 in G[city] else -float("inf")
            return max([dp(week-1, o) + days[city][week] for o in G[city]])
        
        return max([dp(k, city) for city in range(n)])