Problem: Find Number of Coins to Place in Tree Nodes

You are given an undirected tree with `n` nodes labeled from `0` to `n - 1`,
and rooted at node `0`. You are given a 2D integer array `edges` of length `n
- 1`, where `edges[i] = [ai, bi]` indicates that there is an edge between
nodes `ai` and `bi` in the tree.

You are also given a 0-indexed integer array `cost` of length `n`, where
`cost[i]` is the cost assigned to the `ith` node.

You need to place some coins on every node of the tree. The number of coins to
be placed at node `i` can be calculated as:

* If size of the subtree of node `i` is less than `3`, place `1` coin.
* Otherwise, place an amount of coins equal to the maximum product of cost values assigned to `3` distinct nodes in the subtree of node `i`. If this product is negative, place `0` coins.

Return an array `coin` of size `n` such that `coin[i]` is the number of coins
placed at node `i`.


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator

class ChildCost:
  def __init__(self, cost: int):
    self.numNodes = 1
    self.maxPosCosts = [cost] if cost > 0 else []
    self.minNegCosts = [cost] if cost < 0 else []

  def update(self, childCost: 'ChildCost') -> None:
    self.numNodes += childCost.numNodes
    self.maxPosCosts.extend(childCost.maxPosCosts)
    self.minNegCosts.extend(childCost.minNegCosts)
    self.maxPosCosts.sort(reverse=True)
    self.minNegCosts.sort()
    self.maxPosCosts = self.maxPosCosts[:3]
    self.minNegCosts = self.minNegCosts[:2]

  def maxProduct(self) -> int:
    if self.numNodes < 3:
      return 1
    if not self.maxPosCosts:
      return 0
    res = 0
    if len(self.maxPosCosts) == 3:
      res = self.maxPosCosts[0] * self.maxPosCosts[1] * self.maxPosCosts[2]
    if len(self.minNegCosts) == 2:
      res = max(res, self.minNegCosts[0] * self.minNegCosts[1] * self.maxPosCosts[0])
    return res


class Solution:
  def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
    n = len(cost)
    ans = [0] * n
    tree = [[] for _ in range(n)]

    for u, v in edges:
      tree[u].append(v)
      tree[v].append(u)

    def dfs(u: int, prev: int) -> None:
      res = ChildCost(cost[u])
      for v in tree[u]:
        if v != prev:
          res.update(dfs(v, u))
      ans[u] = res.maxProduct()
      return res

    dfs(0, -1)
    return ans


Java solution:
class ChildCost {
  public ChildCost(int cost) {
    if (cost > 0)
      maxPosCosts.add(cost);
    else
      minNegCosts.add(cost);
  }

  public void update(ChildCost childCost) {
    numNodes += childCost.numNodes;
    maxPosCosts.addAll(childCost.maxPosCosts);
    minNegCosts.addAll(childCost.minNegCosts);
    maxPosCosts.sort(Comparator.reverseOrder());
    minNegCosts.sort(Comparator.naturalOrder());
    if (maxPosCosts.size() > 3)
      maxPosCosts = maxPosCosts.subList(0, 3);
    if (minNegCosts.size() > 2)
      minNegCosts = minNegCosts.subList(0, 2);
  }

  public long maxProduct() {
    if (numNodes < 3)
      return 1;
    if (maxPosCosts.isEmpty())
      return 0;
    long res = 0;
    if (maxPosCosts.size() == 3)
      res = (long) maxPosCosts.get(0) * maxPosCosts.get(1) * maxPosCosts.get(2);
    if (minNegCosts.size() == 2)
      res = Math.max(res, (long) minNegCosts.get(0) * minNegCosts.get(1) * maxPosCosts.get(0));
    return res;
  }

  private int numNodes = 1;
  private List<Integer> maxPosCosts = new ArrayList<>();
  private List<Integer> minNegCosts = new ArrayList<>();
}

class Solution {
  public long[] placedCoins(int[][] edges, int[] cost) {
    final int n = cost.length;
    long[] ans = new long[n];
    List<Integer>[] tree = new List[n];

    for (int i = 0; i < n; i++)
      tree[i] = new ArrayList<>();

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      tree[u].add(v);
      tree[v].add(u);
    }

    dfs(tree, 0, /*prev=*/-1, cost, ans);
    return ans;
  }

  private ChildCost dfs(List<Integer>[] tree, int u, int prev, int[] cost, long[] ans) {
    ChildCost res = new ChildCost(cost[u]);
    for (final int v : tree[u])
      if (v != prev)
        res.update(dfs(tree, v, u, cost, ans));
    ans[u] = res.maxProduct();
    return res;
  }
}


C++ solution:
class ChildCost {
 public:
  ChildCost(int cost) {
    numNodes = 1;
    if (cost > 0)
      maxPosCosts.push_back(cost);
    else
      minNegCosts.push_back(cost);
  }

  void update(ChildCost childCost) {
    numNodes += childCost.numNodes;
    ranges::copy(childCost.maxPosCosts, back_inserter(maxPosCosts));
    ranges::copy(childCost.minNegCosts, back_inserter(minNegCosts));
    ranges::sort(maxPosCosts, greater<int>());
    ranges::sort(minNegCosts);
    maxPosCosts.resize(min(static_cast<int>(maxPosCosts.size()), 3));
    minNegCosts.resize(min(static_cast<int>(minNegCosts.size()), 2));
  }

  long maxProduct() {
    if (numNodes < 3)
      return 1;
    if (maxPosCosts.empty())
      return 0;
    long res = 0;
    if (maxPosCosts.size() == 3)
      res = static_cast<long>(maxPosCosts[0]) * maxPosCosts[1] * maxPosCosts[2];
    if (minNegCosts.size() == 2)
      res = max(res, static_cast<long>(minNegCosts[0]) * minNegCosts[1] *
                         maxPosCosts[0]);
    return res;
  }

 private:
  int numNodes;
  vector<int> maxPosCosts;
  vector<int> minNegCosts;
};

class Solution {
 public:
  vector<long long> placedCoins(vector<vector<int>>& edges, vector<int>& cost) {
    const int n = cost.size();
    vector<long long> ans(n);
    vector<vector<int>> tree(n);

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      tree[u].push_back(v);
      tree[v].push_back(u);
    }

    dfs(tree, 0, /*prev=*/-1, cost, ans);
    return ans;
  }

 private:
  ChildCost dfs(const vector<vector<int>>& tree, int u, int prev,
                const vector<int>& cost, vector<long long>& ans) {
    ChildCost res(cost[u]);
    for (const int v : tree[u])
      if (v != prev)
        res.update(dfs(tree, v, u, cost, ans));
    ans[u] = res.maxProduct();
    return res;
  }
};