BFF.py

#!/usr/bin/env python3
from collections import defaultdict
import copy


def read(fd,t=int):
    return t(fd.readline())


def readlist(fd, t=int):
    return list(map(t, fd.readline().split()))


def solve(N, L):
    L = [0] + L
    twos = set()
    # print([i for i in range(1,N+1)])
    # print([i for i in L[1:]])
    L_r = defaultdict(list )
    for i in range(1,N+1):
        L_r [L[i] ] += [i]
        if i == L[L[i]]:
            twos.add(i)

    # print("two bffs {}".format(twos) )
    # dfs
    P = {}
    maxp  = 0
    for i in range(1,N+1):

        if i in twos:
            P[i] = set([i, L[i]])
            continue
        else:
            path = set([i])
            nxt = L[i]

            while nxt not in path:
                path.add(nxt)
                nxt = L[nxt]
                # Terminate on cycle len = 2
                if nxt in twos:
                    path.add(nxt)
                    path.add(L[nxt])
                    break
                # Terminate on Cycle len > 2
                elif nxt in path :
                    if nxt is not i:
                        path = set()
                        break

            maxp = max(maxp,len(path))
            P[i] = path
    # print(P )
    # Reduce 2 cycle sets, just keep one
    two2 = { a: L[a] for a in twos}
    for a in copy.deepcopy(two2):
        if L[a] < a:
            del two2[a]
    twos = set(two2.keys())
    # iter through cycle nodes, find new max path if possible
    # print("Old Max: {}".format(maxp))
    for node in twos:
        for i in range(1,N+1):
            for j in range(i+1,N+1):
                if node in P[i] and node in P[j]:
                    if P[i] & P[j] == set([node, L[node]]):
                        # P[i] -> node <-> node' <- P[j]
                        #
                        # P[i] -> node <-> node'
                        #         ^
                        # P[j] ---|
                        # Todo: Fix condition there for correct checking on Case shown above
                        
                        if node in L_r[i] and node in L_r[j] or node in L_r[j] and node not in L_r[i]:
                            pass
                        else:
                            maxp = max(len(P[i] | P[j]),maxp)
                    else:
                        # print("sth wrong {}, {} {}, {}", P[i],P[j], set([node, L[node]]),P[i] & P[j]  )
                        pass
    return maxp


with open("C-small-practice.in") as f:
    T= read(f)
    for t in range(T):
        N = read(f)
        L = readlist(f)
        ans = solve(N,L)
        print("Case #{}: {}".format(t+1,ans))