We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
  1. Prepare
  2. Algorithms
  3. Search
  4. Cut the Tree
  5. Discussions

Cut the Tree

  • mss2518
     + 1 comment

    We can make a mapping of each node, to the sum of their tree. In this case, the root node has a sum equal to the entire tree, and each other node the respective cost of their subtree.

    If we cut at an edge, we spawn two subtrees: one where we just cut (with a sum equal to the mapping of that node) and a 'bigger' tree, with sum equal to the entire tree (i.e., mapping[1]) - the tree we just removed. With this idea, we can map each node to the sum of its tree in O(V + E), and then cut every edge once, perfoming the check in O(1), giving a further O(V + E).

    The total time complexity then becomes O(V + E), with O(V + E) space, too.

    def cutTheTree(data, edges):
    # Write your code here
    adj = collections.defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
        
    lookup = {}
    
    def treeSum(node, par):
        if node is None:
            return 0
        if node in lookup:
            return lookup[node]
        ans = data[node - 1]
        for ch in adj[node]:
            if ch == par:
                continue
            ans += treeSum(ch, node)
        lookup[node] = ans
        return ans
    
    treeSum(1, 1)
    minDiff = float("inf")
    def minPart(node, par):
        nonlocal minDiff
        if node is None:
            return
        for ch in adj[node]:
            if ch == par:
                continue
            # cut at child
            tree1 = lookup[ch]
            tree2 = lookup[1] - tree1
            diff = abs(tree1 - tree2)
            minDiff = min(minDiff, diff)
            minPart(ch, node)
    minPart(1, 1)
    return minDiff