Find Number of Coins to Place in Tree Nodes
Source: leetcode
Difficulty: Hard
Tags: dynamic-programming, tree, depth-first-search, sorting, heap-priority-queue

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`*.*


 


**Example 1:**


![](https://assets.leetcode.com/uploads/2023/11/09/screenshot-2023-11-10-012641.png)

```

Input: edges = [[0,1],[0,2],[0,3],[0,4],[0,5]], cost = [1,2,3,4,5,6]
Output: [120,1,1,1,1,1]
Explanation: For node 0 place 6 * 5 * 4 = 120 coins. All other nodes are leaves with subtree of size 1, place 1 coin on each of them.

```

**Example 2:**


![](https://assets.leetcode.com/uploads/2023/11/09/screenshot-2023-11-10-012614.png)

```

Input: edges = [[0,1],[0,2],[1,3],[1,4],[1,5],[2,6],[2,7],[2,8]], cost = [1,4,2,3,5,7,8,-4,2]
Output: [280,140,32,1,1,1,1,1,1]
Explanation: The coins placed on each node are:
- Place 8 * 7 * 5 = 280 coins on node 0.
- Place 7 * 5 * 4 = 140 coins on node 1.
- Place 8 * 2 * 2 = 32 coins on node 2.
- All other nodes are leaves with subtree of size 1, place 1 coin on each of them.

```

**Example 3:**


![](https://assets.leetcode.com/uploads/2023/11/09/screenshot-2023-11-10-012513.png)

```

Input: edges = [[0,1],[0,2]], cost = [1,2,-2]
Output: [0,1,1]
Explanation: Node 1 and 2 are leaves with subtree of size 1, place 1 coin on each of them. For node 0 the only possible product of cost is 2 * 1 * -2 = -4. Hence place 0 coins on node 0.

```

 


**Constraints:**


- `2 <= n <= 2 * 104`
- `edges.length == n - 1`
- `edges[i].length == 2`
- `0 <= ai, bi < n`
- `cost.length == n`
- `1 <= |cost[i]| <= 104`
- The input is generated such that `edges` represents a valid tree.




C++ solution:
using i64 = long long;

constexpr int MAX = 20010;

int cnt[MAX];
i64 c[MAX], dp[MAX][2][4], buf[8];
vector<int> adj[MAX];

void Combine(const i64* a, const i64* b, i64* c) {
  std::copy(a, a + 3, buf);
  std::copy(b, b + 3, buf + 3);
  std::sort(buf, buf + 6, std::greater<i64>());
  std::copy(buf, buf + 3, c);
}

void DFS(int x, int p) {
  cnt[x] = 1;
  dp[x][0][0] = c[x] < 0 ? -1 : c[x];
  dp[x][0][1] = -1;
  dp[x][0][2] = -1;
  dp[x][1][0] = c[x] < 0 ? -c[x] : -1;
  dp[x][1][1] = -1;
  dp[x][1][2] = -1;
  for (int y : adj[x]) {
    if (y == p) {
      continue;
    }
    DFS(y, x);
    cnt[x] += cnt[y];
    Combine(dp[x][0], dp[y][0], dp[x][0]);
    Combine(dp[x][1], dp[y][1], dp[x][1]);
  }
}

class Solution {
public:
    vector<long long> placedCoins(vector<vector<int>>& edges, vector<int>& cost) {
      int n = cost.size();
      for (int i = 0; i < n; ++i) {
        c[i] = cost[i];
        std::vector<int>().swap(adj[i]);
      }
      for (const auto& e : edges) {
        int u = e[0], v = e[1];
        adj[u].push_back(v);
        adj[v].push_back(u);
      }
      DFS(0, -1);
      std::vector<i64> ret(n);
      for (int i = 0; i < n; ++i) {
        if (cnt[i] < 3) {
          ret[i] = 1;
        } else {
          ret[i] = 0;
          if (dp[i][0][2] > 0) {
            ret[i] = std::max(ret[i], dp[i][0][0] * dp[i][0][1] * dp[i][0][2]);
          }
          if (dp[i][0][0] > 0 && dp[i][1][1] > 0) {
            ret[i] = std::max(ret[i], dp[i][0][0] * dp[i][1][0] * dp[i][1][1]);
          }
        }
      }
      return ret;
    }
};

Go solution:
func max(a, b int64) int64 {
    if a > b {
        return a
    }
    return b
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}
 
func dfs(v int, graph *map[int][]int, use *[]bool, cost *[]int, ans *[]int64) [][]int64 {
    // We say the vertex can no longer be used.
    (*use)[v] = true
    pol := []int64{}
    otr := []int64{}
    // Add cost[v] to the desired array
    if (*cost)[v] < 0 {
        otr = append(otr, int64((*cost)[v]))
    } else {
        pol = append(pol, int64((*cost)[v]))
    }
    for _, nextV := range (*graph)[v] {
        if !(*use)[nextV] {
            // Define new array otr and pol
            ansNextV := dfs(nextV, graph, use, cost, ans)
            pol = append(pol, ansNextV[0]...)
            sort.Slice(pol, func(i, j int) bool {
                return pol[i] >= pol[j]
            })
            pol = pol[:min(3, len(pol))]
            otr = append(otr, ansNextV[1]...)
            sort.Slice(otr, func(i, j int) bool {
                return otr[i] <= otr[j]
            })
            otr = otr[:min(3, len(otr))]
        } 
    }
    // If lens of otr and pol arrays
    if len(pol) + len(otr) < 3 {
        (*ans)[v] = 1
    } else {
        var mxProiz int64 
        if len(pol) == 0 {
            mxProiz = otr[0] * otr[1] * otr[2]
        } else if len(otr) == 0 {
            mxProiz = pol[0] * pol[1] * pol[2]
        } else if len(otr) == 1 && len(pol) == 2 {
            mxProiz = otr[0] * pol[0] * pol[1]
        } else if len(otr) >= 2 && len(pol) < 3 {
            mxProiz = otr[0] * otr[1] * pol[0]
        } else if len(otr) == 1 && len(pol) >= 3 {
            mxProiz = pol[0] * pol[1] * pol[2]
        } else {
            potentialMxProiz1 := otr[0] * otr[1] * pol[0]
            potentialMxProiz2 := pol[0] * pol[1] * pol[2]
            mxProiz = max(potentialMxProiz1, potentialMxProiz2)
        }
        (*ans)[v] = mxProiz
    }
    // Define the answer for ans[v]
    if (*ans)[v] < 0 {
        (*ans)[v] = 0
    }
    return [][]int64{pol, otr}
}

func placedCoins(edges [][]int, cost []int) []int64 {
    n := len(cost)
    ans := make([]int64, n)
    // Build the graph
    graph := make(map[int][]int) 
    for _, egde := range edges {
        graph[egde[0]] = append(graph[egde[0]], egde[1])
        graph[egde[1]] = append(graph[egde[1]], egde[0])
    }
    use := make([]bool, n)
    // Start DFS
    dfs(0, &graph, &use, &cost, &ans)
    return ans
}

Java solution:
class X {
    // number of nodes in subtree
    public int numNodes;
    // posMax: [maximum, second maximum, third maximum]
    // negMin: [minimum, second minimum]
    public ArrayList<Integer> posMax, negMin;

    // default constructor
    public X() {
        numNodes = 0;
        posMax = new ArrayList<>();
        negMin = new ArrayList<>();
    }

    // constructor with one cost value (or one tree node)
    public X(int val) {
        numNodes = 1;
        posMax = new ArrayList<>();
        negMin = new ArrayList<>();
        if (val > 0) {
            posMax.add(val);
        } else {
            negMin.add(val);
        }
    }

    // update current node with child's information
    public void update(X x) {
        // add child's nodes
        numNodes += x.numNodes;

        // add child's maximum(s) and minimum(s)
        posMax.addAll(x.posMax);
        negMin.addAll(x.negMin);

        // sort them and select top 3 maximum and top 2 minimum
        posMax.sort(Collections.reverseOrder());
        negMin.sort(null);
        while (posMax.size() > 3) {
            posMax.remove(posMax.size() - 1);
        }
        while (negMin.size() > 2) {
            negMin.remove(negMin.size() - 1);
        }
    }

    // maximum product of subtree
    public long product() {
        // case 1
        if (numNodes < 3) {
            return 1;
        }
        // case 2
        long result = 0;
        // choice 1 (all positive)
        if (posMax.size() == 3) {
            result = (long)posMax.get(0) * posMax.get(1) * posMax.get(2);
        }
        // choice 2 (2 negative 1 positive)
        if (negMin.size() == 2 && !posMax.isEmpty()) {
            result = Math.max(result, (long)negMin.get(0) * negMin.get(1) * posMax.get(0));
        }
        return result;
    }
}

class Solution {
    // Depth-First Search on tree
    // fills coins array
    X fillCoins(int n, ArrayList<ArrayList<Integer>> graph, int[] cost, int curr, int prev, long[] coins) {
        // maintain information of current subtree
        X x = new X(cost[curr]);
        for (int next : graph.get(curr)) {
            // avoid calling function back to parent
            if (next != prev) {
                // fillCoins() returns information of child's subtree
                // update current subtree with child's information
                x.update(fillCoins(n, graph, cost, next, curr, coins));
            }
        }
        // once information of current subtree is completed
        // calculate maximum product
        coins[curr] = x.product();
        return x;
    }

    public long[] placedCoins(int[][] edges, int[] cost) {
        int n = cost.length;
        // make tree in adjacency list form
        ArrayList<ArrayList<Integer>> graph = new ArrayList<>(n);
        for (int i = 0; i < n; i++) {
            graph.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            graph.get(edge[0]).add(edge[1]);
            graph.get(edge[1]).add(edge[0]);
        }
        long[] coins = new long[n];
        // call DFS and fill coins array
        fillCoins(n, graph, cost, 0, -1, coins);
        return coins;
    }
}

JavaScript solution:
var placedCoins = function (edges, cost) {
  var graph = new Map();
  var res = Array(cost.length).fill(0);

  var fill = (u, v) => {
    if (graph.has(u)) graph.get(u).push(v);
    else graph.set(u, [v]);
  };

  for (var [u, v] of edges) {
    fill(u, v);
    fill(v, u);
  }

  var dfs = (node, parent) => {
    var arr = graph.get(node);
    var positives = [];
    var negatives = [];

    if (cost[node] > 0) positives.push(cost[node]);
    if (cost[node] < 0) negatives.push(cost[node]);

    arr.forEach((x) => {
      if (x !== parent) {
        var [p, n] = dfs(x, node);
        positives.push(...p);
        negatives.push(...n);
      }
    });

    positives.sort((a, b) => b - a);
    negatives.sort((a, b) => a - b);

    if (positives.length + negatives.length < 3) res[node] = 1;
    else {
      var val = 0;
      if (positives.length >= 3) {
        positives.length = 3;
        val = positives[0] * positives[1] * positives[2];
      }

      if (negatives.length >= 2) {
        negatives.length = 2;
        if (positives.length) {
          val = Math.max(val, positives[0] * negatives[0] * negatives[1]);
        }
      }

      res[node] = val;
    }

    return [positives, negatives];
  };

  dfs(0, -1);
  return res;
};

Python solution:
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

Ruby solution:
def placed_coins(edges, cost)
  n = cost.size
  adj = Array.new(n) { [] }
  edges.each do |u, v|
    adj[u] << v
    adj[v] << u
  end

  coin = Array.new(n, 1)
  stack = [[0, -1, false]]
  memo = {} # node => [cnt, max_vals, min_vals]

  until stack.empty?
    node, parent, visited = stack.pop
    if !visited
      stack << [node, parent, true]
      adj[node].each do |child|
        stack << [child, node, false] unless child == parent
      end
    else
      cnt = 1
      max_vals = []
      min_vals = []
      current_cost = cost[node]
      if current_cost >= 0
        max_vals << current_cost
      else
        min_vals << -current_cost
      end

      adj[node].each do |child|
        next if child == parent
        child_cnt, child_max, child_min = memo[child]
        cnt += child_cnt
        max_vals += child_max
        min_vals += child_min
      end

      max_vals.sort!.reverse!
      max_vals = max_vals[0...3]

      min_vals.sort!.reverse!
      min_vals = min_vals[0...2]

      if cnt >= 3
        product = 0
        if max_vals.size >= 3
          product1 = max_vals[0] * max_vals[1] * max_vals[2]
          product = product1 if product1 > product
        end
        if min_vals.size >= 2 && max_vals.size >= 1
          product2 = min_vals[0] * min_vals[1] * max_vals[0]
          product = [product, product2].max
        end
        coin[node] = product > 0 ? product : 0
      else
        coin[node] = 1
      end

      memo[node] = [cnt, max_vals, min_vals]
    end
  end

  coin
end