Problem statement

https://binarysearch.com/problems/Country-Roads/

Solution

In fact what is asked is to find parts of bipartite graph (tree), and there is only two ways to do it (A vs B and B vs A). So just run dfs with colors.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, source, dest, P):
        G = defaultdict(list)
        n = len(P)
        for x, y in zip(source, dest):
            G[x] += [y]
            G[y] += [x]

        colors = [-1] * n

        def dfs(node, c):
            colors[node] = c
            for neib in G[node]:
                if colors[neib] != -1: continue
                dfs(neib, 1 - c)

        dfs(0, 0)
        s1 = sum(x*y for x, y in zip(colors, P))
        return max(s1, sum(P) - s1)