Maximum Genetic Difference Query LeetCode solution in Python

1
# Runtime: 6663 ms (Top 46.4%) | Memory: 339.09 MB (Top 7.1%)
2

3

4
class Solution:
5
    def maxGeneticDifference(
6
        self, parents: List[int], queries: List[List[int]]
7
    ) -> List[int]:
8
        adj_map = collections.defaultdict(set)
9
        root = None
10
        for i, node in enumerate(parents):
11
            if node == -1:
12
                root = i
13
            else:
14
                adj_map[node].add(i)
15
        queries_map = collections.defaultdict(set)
16

17
        for q in queries:
18
            queries_map[q[0]].add(q[1])
19
        self.res_map = {}
20

21
        def helperDFS(curr_root, prefix_map):
22

23
            if curr_root in queries_map:
24
                for val in queries_map[curr_root]:
25
                    bin_rep = format(val, "020b")
26
                    best_bin = ""
27
                    print(bin_rep)
28
                    for i in range(20):
29
                        if prefix_map[best_bin + str(1 - int(bin_rep[i]))]:
30
                            best_bin += str(1 - int(bin_rep[i]))
31
                        else:
32
                            best_bin += bin_rep[i]
33
                    self.res_map[(curr_root, val)] = int(best_bin, 2) ^ val
34
            for child in adj_map[curr_root]:
35
                bin_rep = format(child, "020b")
36
                for i in range(1, len(bin_rep) + 1):
37
                    prefix_map[bin_rep[0:i]].add(child)
38
                helperDFS(child, prefix_map)
39
                for i in range(1, len(bin_rep) + 1):
40
                    prefix_map[bin_rep[0:i]].remove(child)
41

42
        initial_prefixmap = collections.defaultdict(set)
43
        root_rep = format(root, "020b")
44
        for i in range(1, 21):
45
            initial_prefixmap[root_rep[0:i]].add(root)
46
        helperDFS(root, initial_prefixmap)
47
        res = []
48
        for q in queries:
49
            res.append(self.res_map[(q[0], q[1])])
50
        return res