Maximum Score of a Node Sequence LeetCode solution in Python

1
# defaultdict is the same as Python's usual dictionary, but if an
2
# element doesn't exist, you can give it a default value to initialize with. 
3
from collections import defaultdict
4
# nlargest(n, l) - returns the n largest values of collection l.
5
from heapq import nlargest
6
# "product" is a function that takes two collections and 
7
# returns every pair between them.
8
# product("ab", "cd") = [(a, c), (a, d), (b, c), (b, d)].
9
from itertools import product
10

11
Let V be the number of nodes and E = len(edges).
12
# Time complexity: O(V + E) - we iterate through every vertex 
13
#                  and every edge a constant number of times.
14
# Space complexity: O(V) - we save a constant 
15
#                   number of neighbors (3) for every node.
16
class Solution:
17
    def maximumScore(self, scores: List[int], edges: List[List[int]]) -> int:
18
        # Turn the edge list into an adjacency graph.
19
        m = defaultdict(list)
20
        for u, v in edges:
21
            m[u].append((scores[v], v)) 
22
            m[v].append((scores[u], u))
23

24
        # Cut down all neighbors of each node to the three
25
        # that have the highest value.
26
        for u in m:
27
            m[u] = nlargest(3, m[u])
28

29
        ret = -1
30
        # Consider each edge to potentially be (B, C) for a quadruplet.
31
        for b, c in edges:
32
            # For every possible A and D in the neighbors of B and C...
33
            for (aWeight, a), (dWeight, d) in product(m[b], m[c]):
34
                # ... If we have no redundant nodes, it's a quadruplet.
35
                # Since it's the highest value quadruplet we could
36
                # possibly make with B and C, this solution is always accurate.
37
                if a not in [b, c] and d not in [b, c] and a != d:
38
                    ret = max(ret, scores[b] + scores[c] + aWeight + dWeight)
39

40
        return ret