Minimum Score After Removals on a Tree LeetCode solution in Python

1
# Runtime: 1682 ms (Top 100.0%) | Memory: 18.40 MB (Top 71.7%)
2

3

4
class Solution:
5
    def minimumScore(self, nums: List[int], edges: List[List[int]]) -> int:
6
        n = len(nums)
7
        graph = [[] for _ in range(n)]
8
        for u, v in edges:
9
            graph[u].append(v)
10
            graph[v].append(u)
11

12
        def fn(u, p):
13
            nonlocal t
14
            score[u] = nums[u]
15
            tin[u] = (t := t + 1)  # time to enter
16
            for v in graph[u]:
17
                if v != p:
18
                    fn(v, u)
19
                    score[u] ^= score[v]
20
            tout[u] = t  # time to exit
21

22
        t = 0
23
        score = [0] * n
24
        tin = [0] * n
25
        tout = [0] * n
26
        fn(0, -1)
27

28
        ans = inf
29
        for u in range(1, n):
30
            for v in range(u + 1, n):
31
                if (
32
                    tin[v] <= tin[u] and tout[v] >= tout[u]
33
                ):  # enter earlier & exit later == parent
34
                    uu = score[u]
35
                    vv = score[v] ^ score[u]
36
                    xx = score[0] ^ score[v]
37
                elif tin[v] >= tin[u] and tout[v] <= tout[u]:
38
                    uu = score[u] ^ score[v]
39
                    vv = score[v]
40
                    xx = score[0] ^ score[u]
41
                else:
42
                    uu = score[u]
43
                    vv = score[v]
44
                    xx = score[0] ^ score[u] ^ score[v]
45
                ans = min(ans, max(uu, vv, xx) - min(uu, vv, xx))
46
        return ans