2973. Find Number of Coins to Place in Tree Nodes

Problem Statement

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:

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:

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:

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.

Intuition

Approach:
While Traversing through all the childNodes
Keep Track of SubTree length and positive and negative vals

Links

https://leetcode.com/problems/find-number-of-coins-to-place-in-tree-nodes/description/

Video Links

https://youtu.be/VmTxOpyPtIE?t=6325

Approach 1:

C++

typedef long long ll;

class GraphNode{
public:
    int subTree;
    vector<int> pos, neg;

    GraphNode(int cost){
        subTree = 1;
        if(cost > 0)    
            pos.push_back(cost);
        else    
            neg.push_back(cost);
    }

    void updateVal(GraphNode child){
        subTree += child.subTree;
        pos.insert(pos.end(), child.pos.begin(), child.pos.end());
        neg.insert(neg.end(), child.neg.begin(), child.neg.end());

        sort(pos.rbegin(), pos.rend());
        sort(neg.begin(), neg.end());

        pos.resize(min((int)pos.size(), 3));
        neg.resize(min((int)neg.size(), 2));
    }

    ll updateAns(){
        if(subTree < 3){
            return 1;
        }

        ll ans = 0;
        if(pos.size() == 3)
            ans = 1ll*pos[0]*pos[1]*pos[2];
        if(neg.size() == 2 and pos.size() > 0)
            ans = max(ans, 1ll*neg[0]*neg[1]*pos[0]);

        return ans;
    }
};

class Solution {
public:
    GraphNode dfs(int idx, int parent, vector<vector<int>> &graph, vector<int>& cost,vector<ll> &coins){
        GraphNode node = GraphNode(cost[idx]);

        for(auto &childNode: graph[idx]){
            if(childNode != parent){
                GraphNode child = dfs(childNode, idx, graph, cost, coins);
                node.updateVal(child);
            }
        }
        coins[idx] = node.updateAns();
        return node;
    }

    vector<long long> placedCoins(vector<vector<int>>& edges, vector<int>& cost) {
        int n = cost.size();
        vector<vector<int>> graph(n);

        for(auto &it: edges){
            int u = it[0], v = it[1];
            graph[u].push_back(v);
            graph[v].push_back(u);
        }
        vector<ll> coins(n);
        dfs(0, -1, graph, cost, coins);
        
        return coins;
    }
};

Approach 2:

C++

Approach 3:

C++

Approach 4:

C++

Similar Problems