| def get_sparse_neighbor(p: int, n: int, m: int): | |
| """Returns a dictionnary, where the keys are index of 4-neighbor of `p` in the sparse matrix, | |
| and values are tuples (i, j, x), where `i`, `j` are index of neighbor in the normal matrix, | |
| and x is the direction of neighbor. | |
| Arguments: | |
| p {int} -- index in the sparse matrix. | |
| n {int} -- number of rows in the original matrix (non sparse). | |
| m {int} -- number of columns in the original matrix. | |
| Returns: | |
| dict -- dictionnary containing indices of 4-neighbors of `p`. | |
| """ | |
| i, j = p // m, p % m | |
| d = {} | |
| if i - 1 >= 0: | |
| d[(i - 1) * m + j] = (i - 1, j, 0) | |
| if i + 1 < n: | |
| d[(i + 1) * m + j] = (i + 1, j, 0) | |
| if j - 1 >= 0: | |
| d[i * m + j - 1] = (i, j - 1, 1) | |
| if j + 1 < m: | |
| d[i * m + j + 1] = (i, j + 1, 1) | |
| return d | |