Contain Virus LeetCode solution in Python

1
class Solution:
2
    def containVirus(self, mat: List[List[int]]) -> int:
3
        m, n = len(mat), len(mat[0])
4

5
        def dfs(
6
            i, j, visited, nextInfected
7
        ):  # return no. of walls require to quarantined dfs area
8
            if 0 <= i < m and 0 <= j < n and (i, j) not in visited:
9
                if mat[i][j] == 2:  # Already quarantined cell
10
                    return 0
11
                if mat[i][j] == 0:
12
                    nextInfected.add((i, j))  # add cell which will be infected next day
13
                    return 1  # require one wall to quarantined cell from one side
14

15
                else:
16
                    visited.add((i, j))
17
                    return (
18
                        dfs(i - 1, j, visited, nextInfected)
19
                        + dfs(i + 1, j, visited, nextInfected)
20
                        + dfs(i, j - 1, visited, nextInfected)
21
                        + dfs(i, j + 1, visited, nextInfected)
22
                    )  # traverse all four direction
23
            else:
24
                return 0
25

26
        ans = 0
27
        while (
28
            True
29
        ):  # this loop running "how many days we should installing the walls" times
30
            # For every day check which area infect more cells
31
            visited = set()  # Using in dfs
32
            All_nextinfect = set()
33
            stop, walls = (
34
                set(),
35
                0,
36
            )  # here stop store the indices of maximum no. of cells in which we stop spreading of virus this day
37

38
            for i in range(m):
39
                for j in range(n):
40
                    if mat[i][j] == 1 and (i, j) not in visited:
41
                        nextInfected = set()
42
                        a = dfs(i, j, visited, nextInfected)
43

44
                        if len(stop) < len(nextInfected):
45
                            All_nextinfect = (
46
                                All_nextinfect | stop
47
                            )  # leave previous saved area from virus
48
                            stop = nextInfected  # pick new area which we want to save
49
                            walls = a  # require walls
50
                            p, q = i, j  # starting position(indices) of this area
51
                        else:
52
                            All_nextinfect = All_nextinfect | nextInfected
53

54
            if not stop:  # if our job is done i.e. No cell will be infect Later
55
                break
56
            ans += walls  # add new walls installed this day
57

58
            # change each cell value to 2 which will be covered by quarantined area
59
            def fun(p, q):
60
                if 0 <= p < m and 0 <= q < n and mat[p][q] == 1:
61
                    mat[p][q] = 2
62
                    fun(p + 1, q)
63
                    fun(p - 1, q)
64
                    fun(p, q - 1)
65
                    fun(p, q + 1)
66

67
            fun(p, q)  # start dfs from start point of quarantined area
68

69
            for (
70
                a,
71
                b,
72
            ) in (
73
                All_nextinfect
74
            ):  # set new infected cell value = 1 for iterating next day
75
                mat[a][b] = 1
76

77
        return ans  # Final answer