|
| 1 | +# https://leetcode.com/problems/maximize-the-number-of-target-nodes-after-connecting-trees-i/description/ |
| 2 | + |
| 3 | +def build_graph(edges, n): |
| 4 | + graph = [[] for _ in range(n)] |
| 5 | + for a, b in edges: |
| 6 | + graph[a].append(b) |
| 7 | + graph[b].append(a) |
| 8 | + return graph |
| 9 | + |
| 10 | +def dfs(graph, node, parent, depth): |
| 11 | + if depth < 0: |
| 12 | + return 0 |
| 13 | + count = 1 # Đếm chính nó |
| 14 | + for neighbor in graph[node]: |
| 15 | + if neighbor != parent: |
| 16 | + count += dfs(graph, neighbor, node, depth - 1) |
| 17 | + return count |
| 18 | + |
| 19 | +def max_target_nodes(edges1, edges2, k): |
| 20 | + n1 = len(edges1) + 1 |
| 21 | + n2 = len(edges2) + 1 |
| 22 | + graph1 = build_graph(edges1, n1) |
| 23 | + graph2 = build_graph(edges2, n2) |
| 24 | + |
| 25 | + # Bước 1: Tìm số lượng lớn nhất các nút có thể reach được trong cây 2 với (k - 1) bước |
| 26 | + max_targets_tree2 = 0 |
| 27 | + for i in range(n2): |
| 28 | + reachable = dfs(graph2, i, -1, k - 1) |
| 29 | + max_targets_tree2 = max(max_targets_tree2, reachable) |
| 30 | + |
| 31 | + # Bước 2: Với mỗi nút trong cây 1, tính tổng số target nodes |
| 32 | + result = [] |
| 33 | + for i in range(n1): |
| 34 | + reachable_in_tree1 = dfs(graph1, i, -1, k) |
| 35 | + total = reachable_in_tree1 + max_targets_tree2 |
| 36 | + result.append(total) |
| 37 | + return result |
| 38 | + |
| 39 | +# === MAIN TEST CASE === |
| 40 | +if __name__ == "__main__": |
| 41 | + edges1 = [[0,1],[0,2],[2,3],[2,4]] |
| 42 | + edges2 = [[0,1],[0,2],[0,3],[2,7],[1,4],[4,5],[4,6]] |
| 43 | + k = 2 |
| 44 | + |
| 45 | + output = max_target_nodes(edges1, edges2, k) |
| 46 | + print("Output:", output) # Expected: [9, 7, 9, 8, 8] |
0 commit comments