|
| 1 | +#!/bin/python3 |
| 2 | + |
| 3 | +import math |
| 4 | +import os |
| 5 | + |
| 6 | +class TreeNode: |
| 7 | + def __init__(self, value, children): |
| 8 | + self.value = value |
| 9 | + self.children = children |
| 10 | + self.total_sum = None |
| 11 | + |
| 12 | + def __repr__(self): |
| 13 | + return "TreeNode(%s, %s)" % (self.value, self.total_sum) |
| 14 | + |
| 15 | +def build_tree(tree_values, tree_edges): |
| 16 | + tree_nodes = [TreeNode(v, set()) for v in tree_values] |
| 17 | + for node_from, node_to in tree_edges: |
| 18 | + # The tree input is undirected so I am adding both as children and parent |
| 19 | + # I am later cleaning it up while doing DFS over the tree |
| 20 | + tree_nodes[node_from - 1].children.add(tree_nodes[node_to - 1]) |
| 21 | + tree_nodes[node_to - 1].children.add(tree_nodes[node_from - 1]) |
| 22 | + return tree_nodes[0] |
| 23 | + |
| 24 | +def is_even_number(value): |
| 25 | + return not value & 1 |
| 26 | + |
| 27 | +def populate_tree_sums(root): |
| 28 | + stack = (root, None) |
| 29 | + visited = set() |
| 30 | + |
| 31 | + while stack: |
| 32 | + selected_node = stack[0] |
| 33 | + |
| 34 | + if selected_node not in visited: |
| 35 | + visited.add(selected_node) |
| 36 | + for child in selected_node.children: |
| 37 | + #remove non children, this cleans out the "bad" inputs |
| 38 | + #the tree has undirected edges so when we convert it back to a |
| 39 | + #proper tree it's easier to work with... |
| 40 | + child.children.remove(selected_node) |
| 41 | + #populate the stack: |
| 42 | + stack = (child, stack) |
| 43 | + else: |
| 44 | + stack = stack[-1] # pop stack |
| 45 | + #calculate the total sum of the current node before going up the tree |
| 46 | + selected_node.total_sum = sum( |
| 47 | + map( |
| 48 | + lambda tn: tn.total_sum, |
| 49 | + selected_node.children |
| 50 | + ) |
| 51 | + ) + selected_node.value |
| 52 | + |
| 53 | +def find_best_balanced_forest(root): |
| 54 | + stack = (root, None) |
| 55 | + #visited - visited nodes |
| 56 | + #visited_sums - sums that are currently visited |
| 57 | + #root complement sums complement sums (total_value - parent) on the way to the |
| 58 | + #current node, the cardinality of root complement sums is increased when going |
| 59 | + #down the tree and decreased when going up the tree, it is okay to do that |
| 60 | + #because the sums are always unique in the root_complement_sums |
| 61 | + visited, visited_sums, root_complement_sums = set(), set(), set() |
| 62 | + min_result_value = math.inf |
| 63 | + |
| 64 | + while stack: |
| 65 | + selected_node = stack[0] |
| 66 | + |
| 67 | + if selected_node not in visited: |
| 68 | + visited.add(selected_node) |
| 69 | + |
| 70 | + #populate stack with children all at once: |
| 71 | + for child in selected_node.children: |
| 72 | + stack = (child, stack) |
| 73 | + |
| 74 | + #this is a complement sum: TOTAL - current_sum |
| 75 | + #I need to calculate it while going down the tree so when I go up |
| 76 | + #I will use those values in the root_complement_sums to check for |
| 77 | + #existance |
| 78 | + selected_sum_comp = root.total_sum - selected_node.total_sum |
| 79 | + root_complement_sums.add(selected_sum_comp) |
| 80 | + |
| 81 | + # Yes, no bitwise shifts, I present what I want to get accomplished, |
| 82 | + # but I don't care how it is accomplished |
| 83 | + # selected_node.total_sum * 3 >= root.total_sum is checking that |
| 84 | + # that if the cut is made in selected subtree and the visited subtree |
| 85 | + # (in case the comp or sum exists in the visited sums) |
| 86 | + # the remaining subtree sum is equal or less than the sums |
| 87 | + # (which are equal) of the current and the visited subtrees |
| 88 | + # this is just part of the requirement - I can balance the remaining |
| 89 | + # tree only with 0 or positive elements |
| 90 | + if ( |
| 91 | + (selected_node.total_sum * 2) in visited_sums or |
| 92 | + (root.total_sum - selected_node.total_sum * 2) in visited_sums |
| 93 | + ) and selected_node.total_sum * 3 >= root.total_sum: |
| 94 | + |
| 95 | + #get the candidate value and update min_result_value if it's less |
| 96 | + candidate_value = selected_node.total_sum * 3 - root.total_sum |
| 97 | + if candidate_value < min_result_value: |
| 98 | + min_result_value = candidate_value |
| 99 | + else: |
| 100 | + # This is a case where two even halfs are found. |
| 101 | + if (selected_node.total_sum * 2) == root.total_sum: |
| 102 | + candidate_value = selected_node.total_sum |
| 103 | + # In this case a balanced forest is these two halfs + a new node as |
| 104 | + # a separate tree with the same value as the half of the existing |
| 105 | + # tree sum |
| 106 | + if candidate_value < min_result_value: |
| 107 | + min_result_value = candidate_value |
| 108 | + |
| 109 | + # check visited sums and root complements |
| 110 | + # root complements are the sums on the way from root to the selected |
| 111 | + # nodes taken from it's parents of if we have a tree |
| 112 | + # (1) |
| 113 | + # / | \ |
| 114 | + # / | \ |
| 115 | + # / | \ |
| 116 | + # (2) (3) (4) |
| 117 | + # / \ | /\ |
| 118 | + # (5)(6) (7) (8)(9) |
| 119 | + # |
| 120 | + # |
| 121 | + # If I am at the node 8, I have the {TOTAL - (8).sum, TOTAL - (4).sum } |
| 122 | + # If I am at the node 9, I have the {TOTAL - (9).sum, TOTAL - (4).sum } |
| 123 | + # If I am at the node 2, I have the {TOTAL - (2).sum } |
| 124 | + if ( |
| 125 | + ( |
| 126 | + selected_node.total_sum in visited_sums or |
| 127 | + selected_node.total_sum in root_complement_sums |
| 128 | + ) and selected_node.total_sum * 3 >= root.total_sum |
| 129 | + ): |
| 130 | + # candidate split: |
| 131 | + candidate_value = selected_node.total_sum * 3 - root.total_sum |
| 132 | + if candidate_value < min_result_value: |
| 133 | + min_result_value = candidate_value |
| 134 | + |
| 135 | + selected_sum_comp = root.total_sum - selected_node.total_sum |
| 136 | + if is_even_number(selected_sum_comp): |
| 137 | + #I am not trying to impress anyone with bitwise shifts here: |
| 138 | + selected_sum_comp_half = selected_sum_comp // 2 |
| 139 | + if selected_sum_comp_half > selected_node.total_sum and ( |
| 140 | + selected_sum_comp_half in visited_sums or |
| 141 | + selected_sum_comp_half in root_complement_sums |
| 142 | + ): |
| 143 | + #same candidate value |
| 144 | + candidate_value = selected_sum_comp_half - selected_node.total_sum |
| 145 | + if candidate_value < min_result_value: |
| 146 | + min_result_value = candidate_value |
| 147 | + |
| 148 | + #remove selected complement from root while going up the tree |
| 149 | + root_complement_sums.remove(selected_sum_comp) |
| 150 | + #added to the visited sums while going up the tree |
| 151 | + visited_sums.add(selected_node.total_sum) |
| 152 | + |
| 153 | + #stack pop: |
| 154 | + stack = stack[-1] |
| 155 | + |
| 156 | + if min_result_value == math.inf: |
| 157 | + min_result_value = -1 |
| 158 | + return min_result_value |
| 159 | + |
| 160 | +# Complete the balancedForest function below. |
| 161 | +def balancedForest(tree_values, tree_edges): |
| 162 | + root = build_tree(tree_values, tree_edges) |
| 163 | + populate_tree_sums(root) |
| 164 | + return find_best_balanced_forest(root) |
| 165 | + |
| 166 | + |
| 167 | +if __name__ == '__main__': |
| 168 | + fptr = open(os.environ['OUTPUT_PATH'], 'w') |
| 169 | + |
| 170 | + q = int(input()) |
| 171 | + |
| 172 | + for q_itr in range(q): |
| 173 | + n = int(input()) |
| 174 | + |
| 175 | + c = list(map(int, input().rstrip().split())) |
| 176 | + |
| 177 | + edges = [] |
| 178 | + |
| 179 | + for _ in range(n - 1): |
| 180 | + edges.append(list(map(int, input().rstrip().split()))) |
| 181 | + |
| 182 | + result = balancedForest(c, edges) |
| 183 | + |
| 184 | + fptr.write(str(result) + '\n') |
| 185 | + |
| 186 | + fptr.close() |
0 commit comments