golang 解法, 执行用时: 92 ms, 内存消耗: 8.7 MB, 提交时间: 2024-12-09 09:23:08

func buildTree(edges [][]int, k int) func(int, int, int) int {
	g := make([][]int, len(edges)+1)
	for _, e := range edges {
		x, y := e[0], e[1]
		g[x] = append(g[x], y)
		g[y] = append(g[y], x)
	}

	var dfs func(int, int, int) int
	dfs = func(x, fa, d int) int {
		if d > k {
			return 0
		}
		cnt := 1
		for _, y := range g[x] {
			if y != fa {
				cnt += dfs(y, x, d+1)
			}
		}
		return cnt
	}
	return dfs
}

func maxTargetNodes(edges1, edges2 [][]int, k int) []int {
	max2 := 0
	if k > 0 {
		dfs := buildTree(edges2, k-1) // 注意这里传的是 k-1
		for i := range len(edges2) + 1 {
			max2 = max(max2, dfs(i, -1, 0))
		}
	}

	dfs := buildTree(edges1, k)
	ans := make([]int, len(edges1)+1)
	for i := range ans {
		ans[i] = dfs(i, -1, 0) + max2
	}
	return ans
}

java 解法, 执行用时: 320 ms, 内存消耗: 45 MB, 提交时间: 2024-12-09 09:22:50

class Solution {
    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        int max2 = 0;
        if (k > 0) {
            List<Integer>[] g = buildTree(edges2);
            for (int i = 0; i < edges2.length + 1; i++) {
                max2 = Math.max(max2, dfs(i, -1, 0, g, k - 1)); // 注意这里传的是 k-1
            }
        }

        List<Integer>[] g = buildTree(edges1);
        int[] ans = new int[edges1.length + 1];
        for (int i = 0; i < ans.length; i++) {
            ans[i] = dfs(i, -1, 0, g, k) + max2;
        }
        return ans;
    }

    private List<Integer>[] buildTree(int[][] edges) {
        List<Integer>[] g = new ArrayList[edges.length + 1];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int[] e : edges) {
            int x = e[0];
            int y = e[1];
            g[x].add(y);
            g[y].add(x);
        }
        return g;
    }

    private int dfs(int x, int fa, int d, List<Integer>[] g, int k) {
        if (d > k) {
            return 0;
        }
        int cnt = 1;
        for (int y : g[x]) {
            if (y != fa) {
                cnt += dfs(y, x, d + 1, g, k);
            }
        }
        return cnt;
    }
}

cpp 解法, 执行用时: 143 ms, 内存消耗: 57.8 MB, 提交时间: 2024-12-09 09:22:37

class Solution {
    vector<vector<int>> buildTree(vector<vector<int>>& edges) {
        vector<vector<int>> g(edges.size() + 1);
        for (auto& e : edges) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }
        return g;
    }

    int dfs(int x, int fa, int d, vector<vector<int>>& g, int k) {
        if (d > k) {
            return 0;
        }
        int cnt = 1;
        for (int y : g[x]) {
            if (y != fa) {
                cnt += dfs(y, x, d + 1, g, k);
            }
        }
        return cnt;
    }

public:
    vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2, int k) {
        int max2 = 0;
        if (k > 0) {
            auto g = buildTree(edges2);
            for (int i = 0; i < edges2.size() + 1; i++) {
                max2 = max(max2, dfs(i, -1, 0, g, k - 1)); // 注意这里传的是 k-1
            }
        }

        auto g = buildTree(edges1);
        vector<int> ans(edges1.size() + 1);
        for (int i = 0; i < ans.size(); i++) {
            ans[i] = dfs(i, -1, 0, g, k) + max2;
        }
        return ans;
    }
};

python3 解法, 执行用时: 1655 ms, 内存消耗: 18.4 MB, 提交时间: 2024-12-09 09:22:23

class Solution:
    def buildTree(self, edges: list[list[int]], k: int) -> Callable[[int, int, int], int]:
        g = [[] for _ in range(len(edges) + 1)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)

        def dfs(x: int, fa: int, d: int) -> int:
            if d > k:
                return 0
            cnt = 1
            for y in g[x]:
                if y != fa:
                    cnt += dfs(y, x, d + 1)
            return cnt
        return dfs

    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]], k: int) -> List[int]:
        max2 = 0
        if k:
            dfs = self.buildTree(edges2, k - 1)  # 注意这里传的是 k-1
            max2 = max(dfs(i, -1, 0) for i in range(len(edges2) + 1))

        dfs = self.buildTree(edges1, k)
        return [dfs(i, -1, 0) + max2 for i in range(len(edges1) + 1)]