kitty 27s calculations on a tree hackerrank solution

1from collections import Counter, defaultdict
2
3MOD = 10**9 + 7
4
5def read_row():
6    return (int(x) for x in input().split())
7
8def mul(x, y):
9    return (x * y) % MOD
10
11def add(*args):
12    return sum(args) % MOD
13
14def sub(x, y):
15    return (x - y) % MOD
16
17n, q = read_row()
18
19# Construct adjacency list of the tree
20adj_list = defaultdict(list)
21
22for _ in range(n - 1):
23    u, v = read_row()
24    adj_list[u].append(v)
25    adj_list[v].append(u)
26
27# Construct element to set mapping {element: [sets it belongs to]}
28elements = {v: set() for v in adj_list}
29
30for set_no in range(q):
31    read_row()
32    for x in read_row():
33        elements[x].add(set_no)
34
35# Do BFS to find parent for each node and order them in reverse depth
36root = next(iter(adj_list))
37current = [root]
38current_depth = 0
39order = []
40parent = {root: None}
41depth = {root: current_depth}
42
43while current:
44    current_depth += 1
45    order.extend(current)
46    nxt = []
47    for node in current:
48        for neighbor in adj_list[node]:
49            if neighbor not in parent:
50                parent[neighbor] = node
51                depth[neighbor] = current_depth
52                nxt.append(neighbor)
53
54    current = nxt
55
56# Process nodes in the order created above
57score = Counter()
58# {node: {set_a: [depth, sum of nodes, flow]}}
59state = {}
60for node in reversed(order):
61    states = [state[neighbor] for neighbor in adj_list[node] if neighbor != parent[node]]
62    largest = {s: [depth[node], node, 0] for s in elements[node]}
63
64    if states:
65        max_index = max(range(len(states)), key=lambda x: len(states[x]))
66        if len(states[max_index]) > len(largest):
67            states[max_index], largest = largest, states[max_index]
68
69
70    sets = defaultdict(list)
71    for cur_state in states:
72        for set_no, v in cur_state.items():
73            sets[set_no].append(v)
74
75    for set_no, states in sets.items():
76        if len(states) == 1 and set_no not in largest:
77            largest[set_no] = states[0]
78            continue
79
80        if set_no in largest:
81            states.append(largest.pop(set_no))
82
83        total_flow = 0
84        total_node_sum = 0
85
86        for node_depth, node_sum, node_flow in states:
87            flow_delta = mul(node_depth - depth[node], node_sum)
88            total_flow = add(total_flow, flow_delta, node_flow)
89            total_node_sum += node_sum
90
91        set_score = 0
92
93        for node_depth, node_sum, node_flow in states:
94            node_flow = add(mul(node_depth - depth[node], node_sum), node_flow)
95            diff = mul(sub(total_flow, node_flow), node_sum)
96            set_score = add(set_score, diff)
97
98        score[set_no] = add(score[set_no], set_score)
99        largest[set_no] = (depth[node], total_node_sum, total_flow)
100
101    state[node] = largest
102
103print(*(score[i] for i in range(q)), sep='\n')