Cat and Mouse LeetCode solution in Python

1
class Solution:
2
    def catMouseGame(self, graph: List[List[int]]) -> int:
3
        def getPreStates(m, c, t):
4
            ans = []
5
            if t == 1:
6
                for c2 in graph[c]:
7
                    if c2 == 0:
8
                        continue
9
                    ans.append((m, c2, 2))
10
            else:
11
                for m2 in graph[m]:
12
                    ans.append((m2, c, 1))
13
            return ans
14

15
        def ifAllNextMovesFailed(m, c, t):
16
            if t == 1:
17
                for m2 in graph[m]:
18
                    if result[(m2, c, 2)] != 2:
19
                        return False
20
            else:
21
                for c2 in graph[c]:
22
                    if c2 == 0:
23
                        continue
24
                    if result[(m, c2, 1)] != 1:
25
                        return False
26
            return True
27

28
        result = defaultdict(lambda: 0)
29
        # key = (m,c,turn) value = (0/1/2)
30
        n = len(graph)
31
        queue = deque()
32

33
        for t in range(1, 3):
34
            for i in range(1, n):
35
                # mouse win
36
                result[(0, i, t)] = 1
37
                queue.append((0, i, t))
38
                # cat win
39
                result[(i, i, t)] = 2
40
                queue.append((i, i, t))
41

42
        while queue:
43
            m, c, t = queue.popleft()
44
            r = result[(m, c, t)]
45
            for m2, c2, t2 in getPreStates(m, c, t):
46
                r2 = result[(m2, c2, t2)]
47
                if r2 > 0:
48
                    continue
49
                # populate prestate
50
                if r == 3 - t:  # can always win
51
                    result[(m2, c2, t2)] = r
52
                    queue.append((m2, c2, t2))
53
                elif ifAllNextMovesFailed(m2, c2, t2):
54
                    result[(m2, c2, t2)] = 3 - t2
55
                    queue.append((m2, c2, t2))
56
        return result[(1, 2, 1)]