Medium

Maximize the Number of Target Nodes After Connecting Trees I

LAB

For each node in tree1, connect it to one node in tree2 and return the maximum nodes within distance k.

EXAMPLES

Example 1

Input
{
  "edges1": [
    [
      0,
      1
    ],
    [
      0,
      2
    ]
  ],
  "edges2": [
    [
      0,
      1
    ]
  ],
  "k": 1
}

Output
[
  3,
  2,
  2
]

FUNCTION SHAPE

edges1: intMatrixedges2: intMatrixk: int→intArray

SOLUTION NOTE

For each node i, connect to best j in tree2. Count nodes within k in tree1, k-1 in tree2 (one edge used for connection).

Reveal reference solution +

pythonREFERENCE

def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]], k: int) -> List[int]:
    adjList1, adjList2 = defaultdict(list), defaultdict(list)
    for u, v in edges1:
        adjList1[u].append(v)
        adjList1[v].append(u)
    for u, v in edges2:
        adjList2[u].append(v)
        adjList2[v].append(u)

    @cache
    def dp(i, par, k, tree):
        if k < 0: return 0
        if k == 0: return 1
        res = 1
        adj = adjList1[i] if tree else adjList2[i]
        for nbr in adj:
            if nbr != par:
                res += dp(nbr, i, k - 1, tree)
        return res

    n, m = len(edges1) + 1, len(edges2) + 1
    return [max(dp(i, -1, k, True) + dp(j, -1, k - 1, False) for j in range(m)) for i in range(n)]

TimeO(n × m × k)

SpaceO(n × k)

Open on LeetCode ↗