Skip to content

Instantly share code, notes, and snippets.

@Tokiyomi
Created October 24, 2022 06:15
Show Gist options
  • Save Tokiyomi/fe5f9bde75eaa8917669f70385b89929 to your computer and use it in GitHub Desktop.
Save Tokiyomi/fe5f9bde75eaa8917669f70385b89929 to your computer and use it in GitHub Desktop.
def solveTree():
N = int(input()) # Number of nodes/modules/leaves
F = [0] + [int(f) for f in str(input()).split(' ')] # List of fun factors + an extra 0 at the begining to store the abyss (root value)
P = [int(p) for p in str(input()).split(' ')] # List of pointers
adj = [[] for i in range(N+1)] # Create an array of adjacency lists of all the nodes + 1 list for the abyss adj nodes
for i in range(len(P)): # For each pointer
# Add the adj nodes to the node pointed
adj[P[i]].append(i+1) # append the index + 1 as we added the abyss as a node,so adj[0] are the nodes pointing to the abyss then
curr_fun = 0 # Initialice accumulated fun
for curr_node in reversed(range(N)): # Start from the last node registered (reversed order) for convenience
if len(adj[curr_node])==0: # If the current node has no adj nodes, they are initiators, go to the previous one in the list
continue
# otherwise, it is a parent, so do the following
min_node = min(adj[curr_node], key=F.__getitem__) # from the list of adj nodes of curr node, get the node with the minimum F to be activated first, this is like a resumed loop among all the adj nodes and their Fs
F[curr_node] = max(F[curr_node], F[min_node]) # Compare parent node with the min_node. Update the fun factor of curr_node, conserve the one with the biggest F
curr_fun += sum(F[nu] for nu in adj[curr_node] if nu != min_node) # Sum all the other children nodes Fs of curr_node to the total sum of fun
curr_fun += F[0] # After summing all the resulting fun factors, we are only missing to sum the root value of the tree
return curr_fun
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment