Question

Code

class Solution:
    def maximumInvitations(self, favorite: List[int]) -> int:

        '''
        Two cases:

        1. Closed cycle, last node connect to the first node
        2. Cycle with two node and two edges, with other node pointing towards them 

        '''

        N = len(favorite)
        visit = [False] * N 
        longest_cycle = 0 

        length_2_cycles = []

        for i in range(N):
            if visit[i]:
                continue

            start, cur = i, i
            cur_set = set()

            while not visit[cur]:
                visit[cur] = True
                cur_set.add(cur)
                cur = favorite[cur]

            if cur in cur_set:
                length = len(cur_set)
                while start != cur:
                    length -= 1
                    start = favorite[start]

                if length == 2:
                    length_2_cycles.append([cur, favorite[cur]])
                longest_cycle = max(longest_cycle, length)

        inverted = defaultdict(list)

        for dst, src in enumerate(favorite):
            inverted[src].append(dst)

        # Prove that all two edge cycle subset are disjoint. 
        def bfs(src, parent):
            q = deque([(src, 0)])
            max_length = 0

            while q:
                cur, cur_len = q.popleft()
                if cur == parent:
                    continue

                max_length = max(max_length, cur_len)

                for nei in inverted[cur]:
                    q.append([nei, cur_len + 1])
            return max_length

        total_length = 0
        for n1, n2 in length_2_cycles:
            total_length += bfs(n1, n2) + bfs(n2, n1) + 2

        return max(longest_cycle, total_length)