Problem statement

https://binarysearch.com/problems/Connect-Forest/

Solution

The idea is that we have several components and we need to:

  1. Find diameter and center of each component. Also find radius of component.
  2. Then optimal strategy is to connect centers of these components c0, c1, ..., c_{k-1} in the way where c0 connected with all of the others.
  3. How to choose c0? We need to sort the radii of components in decreasing order and choose the biggest one for c0 and the rest for others.
  4. Then there will be 3 types of pathes: first one is pathes inside components, that is max among self.diams. Second type are pathes between central component c0 and one of the others, we have path of length r0 + ri + 1 in this case. Finally we have paths between ci -> c0 -> cj, then we have ri + rj + 2 length.

Complexity

Time complexity if O(n log n), because we sorted our radii. In fact we do not need to do it, we can just choose the 3 biggest and make it in O(n). Space complexity is O(n).

Code

class Solution:
    def solve(self, G): 
        V, n = set(), len(G)
        self.diams = []
        
        def dfs(node):
            V.add(node)
            maxs = [-1, -1]
            for neib in G[node]:
                if neib in V: continue
                depth = dfs(neib)
                maxs = sorted(maxs + [depth])[-2:]
            self.res = max(self.res, sum(maxs) + 2)
            return maxs[1] + 1
        
        for i in range(n):
            self.res = 0
            if i not in V:
                dfs(i)
                self.diams += [self.res]

        if G == [[]]: return 0
        
        k = len(self.diams)
        self.radii = sorted([(i+1)//2 for i in self.diams])[::-1]
        self.ans = max(self.diams)
        if k >= 2: self.ans = max(self.ans, self.radii[0] + self.radii[1] + 1)
        if k >= 3: self.ans = max(self.ans, self.radii[1] + self.radii[2] + 2)

        return self.ans