File size: 869 Bytes
238eab6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
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
|