golang 解法, 执行用时: 256 ms, 内存消耗: 17.8 MB, 提交时间: 2023-12-25 00:38:37

func placedCoins(edges [][]int, cost []int) []int64 {
	n := len(cost)
	g := make([][]int, n)
	for _, e := range edges {
		x, y := e[0], e[1]
		g[x] = append(g[x], y)
		g[y] = append(g[y], x)
	}

	ans := make([]int64, n)
	var dfs func(int, int) []int
	dfs = func(x, fa int) []int {
		a := []int{cost[x]}
		for _, y := range g[x] {
			if y != fa {
				a = append(a, dfs(y, x)...)
			}
		}

		slices.Sort(a)
		m := len(a)
		if m < 3 {
			ans[x] = 1
		} else {
			ans[x] = int64(max(a[m-3]*a[m-2]*a[m-1], a[0]*a[1]*a[m-1], 0))
		}
		if m > 5 {
			a = append(a[:2], a[m-3:]...)
		}
		return a
	}
	dfs(0, -1)
	return ans
}

java 解法, 执行用时: 74 ms, 内存消耗: 59.5 MB, 提交时间: 2023-12-25 00:38:23

class Solution {
    public long[] placedCoins(int[][] edges, int[] cost) {
        int n = cost.length;
        List<Integer>[] g = new ArrayList[n];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (int[] e : edges) {
            int x = e[0], y = e[1];
            g[x].add(y);
            g[y].add(x);
        }

        long[] ans = new long[n];
        dfs(0, -1, cost, g, ans);
        return ans;
    }

    private List<Integer> dfs(int x, int fa, int[] cost, List<Integer>[] g, long[] ans) {
        List<Integer> a = new ArrayList<>();
        a.add(cost[x]);
        for (int y : g[x]) {
            if (y != fa) {
                a.addAll(dfs(y, x, cost, g, ans));
            }
        }
        Collections.sort(a);
        int m = a.size();
        if (m < 3) {
            ans[x] = 1;
        } else {
            ans[x] = Math.max((long) a.get(m - 3) * a.get(m - 2) * a.get(m - 1),
                    Math.max((long) a.get(0) * a.get(1) * a.get(m - 1), 0));
        }
        if (m > 5) {
            a = List.of(a.get(0), a.get(1), a.get(m - 3), a.get(m - 2), a.get(m - 1));
        }
        return a;
    }
}

cpp 解法, 执行用时: 780 ms, 内存消耗: 238.4 MB, 提交时间: 2023-12-25 00:38:10

class Solution {
public:
    vector<long long> placedCoins(vector<vector<int>> &edges, vector<int> &cost) {
        int n = cost.size();
        vector<vector<int>> g(n);
        for (auto &e: edges) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }

        vector<long long> ans(n, 1);
        function<vector<int>(int, int)> dfs = [&](int x, int fa) -> vector<int> {
            vector<int> a = {cost[x]};
            for (int y: g[x]) {
                if (y != fa) {
                    auto res = dfs(y, x);
                    a.insert(a.end(), res.begin(), res.end());
                }
            }
            sort(a.begin(), a.end());
            int m = a.size();
            if (m >= 3) {
                ans[x] = max({(long long) a[m - 3] * a[m - 2] * a[m - 1], (long long) a[0] * a[1] * a[m - 1], 0LL});
            }
            if (m > 5) {
                a = {a[0], a[1], a[m - 3], a[m - 2], a[m - 1]};
            }
            return a;
        };
        dfs(0, -1);
        return ans;
    }
};

python3 解法, 执行用时: 432 ms, 内存消耗: 36.7 MB, 提交时间: 2023-12-25 00:37:56

class Solution:
    def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
        n = len(cost)
        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)

        ans = [1] * n
        def dfs(x: int, fa: int) -> List[int]:
            a = [cost[x]]
            for y in g[x]:
                if y != fa:
                    a.extend(dfs(y, x))
            a.sort()
            m = len(a)
            if m >= 3:
                ans[x] = max(a[-3] * a[-2] * a[-1], a[0] * a[1] * a[-1], 0)
            if m > 5:
                a = a[:2] + a[-3:]
            return a
        dfs(0, -1)
        return ans