<p>There exists an undirected tree rooted at node <code>0</code> with <code>n</code> nodes labeled from <code>0</code> to <code>n - 1</code>. You are given a 2D <strong>integer</strong> array <code>edges</code> of length <code>n - 1</code>, where <code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>]</code> indicates that there is an edge between nodes <code>a<sub>i</sub></code> and <code>b<sub>i</sub></code> in the tree. You are also given a <strong>0-indexed</strong> array <code>coins</code> of size <code>n</code> where <code>coins[i]</code> indicates the number of coins in the vertex <code>i</code>, and an integer <code>k</code>.</p>

<p>Starting from the root, you have to collect all the coins such that the coins at a node can only be collected if the coins of its ancestors have been already collected.</p>

<p>Coins at <code>node<sub>i</sub></code> can be collected in one of the following ways:</p>

<ul>
	<li>Collect all the coins, but you will get <code>coins[i] - k</code> points. If <code>coins[i] - k</code> is negative then you will lose <code>abs(coins[i] - k)</code> points.</li>
	<li>Collect all the coins, but you will get <code>floor(coins[i] / 2)</code> points. If this way is used, then for all the <code>node<sub>j</sub></code> present in the subtree of <code>node<sub>i</sub></code>, <code>coins[j]</code> will get reduced to <code>floor(coins[j] / 2)</code>.</li>
</ul>

<p>Return <em>the <strong>maximum points</strong> you can get after collecting the coins from <strong>all</strong> the tree nodes.</em></p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2023/09/18/ex1-copy.png" style="width: 60px; height: 316px; padding: 10px; background: rgb(255, 255, 255); border-radius: 0.5rem;" />
<pre>
<strong>Input:</strong> edges = [[0,1],[1,2],[2,3]], coins = [10,10,3,3], k = 5
<strong>Output:</strong> 11                        
<strong>Explanation:</strong> 
Collect all the coins from node 0 using the first way. Total points = 10 - 5 = 5.
Collect all the coins from node 1 using the first way. Total points = 5 + (10 - 5) = 10.
Collect all the coins from node 2 using the second way so coins left at node 3 will be floor(3 / 2) = 1. Total points = 10 + floor(3 / 2) = 11.
Collect all the coins from node 3 using the second way. Total points = 11 + floor(1 / 2) = 11.
It can be shown that the maximum points we can get after collecting coins from all the nodes is 11. 
</pre>

<p><strong class="example">Example 2:</strong></p>
<strong class="example"> <img alt="" src="https://assets.leetcode.com/uploads/2023/09/18/ex2.png" style="width: 140px; height: 147px; padding: 10px; background: #fff; border-radius: .5rem;" /></strong>

<pre>
<strong>Input:</strong> edges = [[0,1],[0,2]], coins = [8,4,4], k = 0
<strong>Output:</strong> 16
<strong>Explanation:</strong> 
Coins will be collected from all the nodes using the first way. Therefore, total points = (8 - 0) + (4 - 0) + (4 - 0) = 16.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>n == coins.length</code></li>
	<li><code>2 &lt;= n &lt;= 10<sup>5</sup></code></li>
	<li><code><font face="monospace">0 &lt;= coins[i] &lt;= 10<sup>4</sup></font></code></li>
	<li><code>edges.length == n - 1</code></li>
	<li><code><font face="monospace">0 &lt;= edges[i][0], edges[i][1] &lt; n</font></code></li>
	<li><code><font face="monospace">0 &lt;= k &lt;= 10<sup>4</sup></font></code></li>
</ul>

from math import inf
from collections import deque

def update_list(data: List[int], idx: int, val: int):
    data.extend([None for _ in range(len(data) - idx + 1)])
    data[idx] = val

def list_get(data: List[int], idx: int):
    return None if idx >= len(data) else data[idx]

class Node:
    def __init__(self, id, coin) -> None:
        self.id = id
        self.coin = coin
        self.neighbors: Set[Node] = set()
        self.parent: Optional[Node] = None
        self.children: Set[Node] = set()
        self.max_point_list: List[int] = list()
    def __hash__(self) -> int:
        return self.id
    def add_neighbor(self, neighbor: 'Node'):
        self.neighbors.add(neighbor)
        neighbor.neighbors.add(self)
    def add_child(self, child: 'Node'):
        self.children.add(child)
        child.parent = self
    def __str__(self) -> str:
        return f"{self.id}"
        return f"{self.id} -> {self.children}"
    def __repr__(self) -> str:
        return self.__str__()


class Solution:
    
    def get_max_point(self, root: Node, halving: int, k: int, max_coin_length: int):
        
        halving = min(halving, max_coin_length)
        
        max_point = list_get(root.max_point_list, halving)
        if max_point is not None:
            return max_point
        
        ret, points_1, points_2 = -inf, -inf, -inf
        root_point = root.coin >> halving

        # strategy 1: take (coin - k) points, keep the descendants the same
        points_1 = root_point - k
        for child in root.children:
            points_1 += self.get_max_point(child, halving, k, max_coin_length)
        ret = max(ret, points_1)
        
        # strategy 2: take (coin // 2) points, half the coin of every descendants
        if k > 0:  # only consider strategy 2 if there is punishment for strategy 1
            points_2 = root_point >> 1
            for child in root.children:
                points_2 += self.get_max_point(child, halving + 1, k, max_coin_length)
            ret = max(ret, points_2)
        
        # update memo
        update_list(root.max_point_list, halving, ret)
        
        # strat = 1 if points_1 > points_2 else 2 if points_2 > points_1 else 0
        # print(f"New entry: {(root, halving)} -> {ret} with strat {strat}")
        # print(f"max_point_list of Node {root.id} = {root.max_point_list}")
        # print()
        
        return ret
    
    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:
        
        # build the graph
        nodes: List[Node] = []
        max_coin = -inf
        for node_id, coin in enumerate(coins):
            nodes.append(Node(node_id, coin))
            max_coin = max(max_coin, coin)
        for node0, node1 in edges:
            nodes[node0].add_neighbor(nodes[node1])
        
        # build the tree from the graph
        parents: deque[Node] = deque()
        parents.append(nodes[0])
        while parents:
            parent = parents.popleft()
            for child in parent.neighbors:
                if child is parent.parent:
                    continue
                parent.add_child(child)
                parents.append(child)
                
        return self.get_max_point(nodes[0], 0, k, max_coin.bit_length())

3864