|
21 | 21 | ========================================= |
22 | 22 | This problem can be solved using DFS or BFS. |
23 | 23 | If 1 is found, find all horizontal or vertical neighbours (1s), and mark them as 0. |
24 | | - Time Complexity: O(N*M) |
25 | | - Space Complexity: O(R) , R = num of rivers |
| 24 | + Time Complexity: O(N*M) |
| 25 | + Space Complexity: O(R) , R = num of rivers |
26 | 26 | ''' |
27 | 27 |
|
28 | 28 |
|
|
31 | 31 | ############ |
32 | 32 |
|
33 | 33 | def river_sizes(matrix): |
34 | | - n = len(matrix) |
35 | | - m = len(matrix[0]) |
36 | | - |
37 | | - results = [] |
38 | | - |
39 | | - for i in range(n): |
40 | | - for j in range(m): |
41 | | - if matrix[i][j] != 0: |
42 | | - # find the river size |
43 | | - size = dfs((i, j), matrix) |
44 | | - |
45 | | - # save the river size |
46 | | - results.append(size) |
47 | | - |
48 | | - return results |
49 | | - |
| 34 | + n = len(matrix) |
| 35 | + m = len(matrix[0]) |
| 36 | + |
| 37 | + results = [] |
| 38 | + |
| 39 | + for i in range(n): |
| 40 | + for j in range(m): |
| 41 | + if matrix[i][j] != 0: |
| 42 | + # find the river size |
| 43 | + size = dfs((i, j), matrix) |
| 44 | + |
| 45 | + # save the river size |
| 46 | + results.append(size) |
| 47 | + |
| 48 | + return results |
| 49 | + |
50 | 50 | def dfs(coord, matrix): |
51 | | - (i, j) = coord |
52 | | - |
53 | | - if i < 0 or j < 0: |
54 | | - # invalid position |
55 | | - return 0 |
56 | | - |
57 | | - n = len(matrix) |
58 | | - m = len(matrix[0]) |
59 | | - |
60 | | - if i == n or j == m: |
61 | | - # invalid position |
62 | | - return 0 |
63 | | - |
64 | | - if matrix[i][j] == 0: |
65 | | - # not a river |
66 | | - return 0 |
67 | | - |
68 | | - # update the matrix, the matrix is passed by reference |
69 | | - matrix[i][j] = 0 |
70 | | - # this position is part of river |
71 | | - size = 1 |
72 | | - |
73 | | - # directions: down, left, up, right |
74 | | - dirs = [(-1, 0), (0, -1), (1, 0), (0, 1)] |
75 | | - |
76 | | - # check all 4 directions |
77 | | - for d in dirs: |
78 | | - size += dfs((i + d[0], j + d[1]), matrix) |
79 | | - |
80 | | - return size |
| 51 | + (i, j) = coord |
| 52 | + |
| 53 | + if i < 0 or j < 0: |
| 54 | + # invalid position |
| 55 | + return 0 |
| 56 | + |
| 57 | + n = len(matrix) |
| 58 | + m = len(matrix[0]) |
| 59 | + |
| 60 | + if i == n or j == m: |
| 61 | + # invalid position |
| 62 | + return 0 |
| 63 | + |
| 64 | + if matrix[i][j] == 0: |
| 65 | + # not a river |
| 66 | + return 0 |
| 67 | + |
| 68 | + # update the matrix, the matrix is passed by reference |
| 69 | + matrix[i][j] = 0 |
| 70 | + # this position is part of river |
| 71 | + size = 1 |
| 72 | + |
| 73 | + # directions: down, left, up, right |
| 74 | + dirs = [(-1, 0), (0, -1), (1, 0), (0, 1)] |
| 75 | + |
| 76 | + # check all 4 directions |
| 77 | + for d in dirs: |
| 78 | + size += dfs((i + d[0], j + d[1]), matrix) |
| 79 | + |
| 80 | + return size |
81 | 81 |
|
82 | 82 |
|
83 | 83 | ########### |
|
0 commit comments