Maximum Compatibility Score Sum LeetCode solution in Python

1
# Runtime: 104 ms (Top 82.18%) | Memory: 13.9 MB (Top 66.34%)
2
import heapq
3
from collections import defaultdict
4

5

6
class Solution:
7
    def maxCompatibilitySum(
8
        self, students: List[List[int]], mentors: List[List[int]]
9
    ) -> int:
10
        m, n = len(students), len(students[0])
11

12
        def hamming(student, mentor):
13
            return sum([int(student[i] != mentor[i]) for i in range(n)])
14

15
        pq = [
16
            (0, 0, "0" * m)
17
        ]  # state: (n-comp_score aka Hamming distance, number of assigned students, mentor status)
18
        optimal = defaultdict(lambda: float("inf"))
19

20
        while pq:  # O(V)
21
            cost, i, mentor_status = heapq.heappop(pq)  # O(logV)
22

23
            # early stopping with termination condition
24
            if i == m:
25
                return m * n - cost
26

27
            # generate successors. The next student to be assigned is at index i
28
            for j, mentor in enumerate(mentors):  # O(m)
29
                if mentor_status[j] != "1":
30
                    new_cost = cost + hamming(students[i], mentor)
31
                    new_mentor_status = mentor_status[:j] + "1" + mentor_status[j + 1 :]
32

33
                    # update optimal cost if a new successor appears with lower cost to the same node
34
                    if new_cost < optimal[(i + 1, new_mentor_status)]:
35
                        optimal[(i + 1, new_mentor_status)] = new_cost
36
                        heapq.heappush(
37
                            pq, (new_cost, i + 1, new_mentor_status)
38
                        )  # O(logV)
39

40
        return 0