Problem of The Day: Maximize the Number of Target Nodes After Connecting Trees II

Problem Statement

leetcode problem link

Brute Force [TLE]

class Solution:
    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]]) -> List[int]:
        n = len(edges1) + 1
        m = len(edges2) + 1
        graph1 = {i: [] for i in range(n)}
        graph2 = {j: [] for j in range(m)}
        nodes1 = {i: 0 for i in range(n)}
        nodes2 = {j: 0 for j in range(m)}
        res = []

        # print(graph1, n)
        # print(graph2, m)

        for u, v in edges1:
            graph1[u].append(v)
            graph1[v].append(u)
        for u, v in edges2:
            graph2[u].append(v)
            graph2[v].append(u)

        def bfs(start, graph, nodes):
            queue = deque([[start, 0]])
            visited = set()
            visited.add(start)
            count = 0
            while queue:
                node, level = queue.popleft()
                if level % 2 == 0:
                    count += 1
                for nei in graph[node]:
                    if nei not in visited:
                        queue.append([nei, level + 1])
                        visited.add(nei)
            nodes[start] = count

        max_nodes2 = 0
        for i in range(n):
            bfs(i, graph1, nodes1)

        for i in range(m):
            bfs(i, graph2, nodes2)
            max_nodes2 = max(max_nodes2, nodes2[i])

        # print(nodes1)
        # print(nodes2)
        # print(max_nodes2)

        for i in range(n):
            res.append(nodes1[i] + max_nodes2)

        return res

Editorial Solution

Approach: Depth-First Search

class Solution:
    def maxTargetNodes(
        self, edges1: List[List[int]], edges2: List[List[int]]
    ) -> List[int]:
        def dfs(node, parent, depth, children, color):
            res = 1 - depth % 2
            color[node] = depth % 2
            for child in children[node]:
                if child == parent:
                    continue
                res += dfs(child, node, depth + 1, children, color)
            return res

        def build(edges, color):
            n = len(edges) + 1
            children = [[] for _ in range(n)]
            for u, v in edges:
                children[u].append(v)
                children[v].append(u)
            res = dfs(0, -1, 0, children, color)
            return [res, n - res]

        n = len(edges1) + 1
        m = len(edges2) + 1
        color1 = [0] * n
        color2 = [0] * m
        count1 = build(edges1, color1)
        count2 = build(edges2, color2)
        res = [0] * n
        for i in range(n):
            res[i] = count1[color1[i]] + max(count2[0], count2[1])
        return res