Problem statement

https://leetcode.com/problems/largest-color-value-in-a-directed-graph/

Solution

First, we use topological sort to check if we have loops. Then for each symbol s, we denote by dp[i] the maximum number of letters s in paths, ending with i. Because we already did topological sort, no need to use memoisation, we can traverse nodes in correct order. For each i we need to check all nodes j from which we can go to i and choose the maximum one plus indicator that C[i] == s.

Complexity

Total time complexity is O(26*(n+m)), where n is number of nodes and m is number of edges. Space complexity is O(n).

Code

class Solution:
    def largestPathValue(self, C, edges):
        n, ans = len(C), 0
        pre, suc = defaultdict(int), defaultdict(list)
        for a, b in edges:
            pre[a] += 1
            suc[b].append(a)
        
        free, out = set(range(n)) - set(pre), []
        while free:
            a = free.pop()
            out.append(a)
            for b in suc[a]:
                pre[b] -= 1
                if not pre[b]: free.add(b)
                    
        if len(out) != n: return -1
        
        for s in set(C):
            res = [0]*n
            for i in out[::-1]:
                res[i] = max([res[j] for j in suc[i]] or [0]) + int(C[i] == s)
            ans = max(ans, max(res))
        
        return ans

We can do the second part using lru_cache. Be careful, if we change order of for loop for s and i, we will get TLE: because we need to jump a lot in indexes.

@lru_cache(None)
    def dp(node, s):
        return max([dp(i, s) for i in suc[node]] or [0]) + int(C[node] == s)  

    for s in set(C):
        for i in range(n):
            ans = max(ans, dp(i, s))
                
    return ans