-
Notifications
You must be signed in to change notification settings - Fork 70
/
utils.py
26 lines (22 loc) · 869 Bytes
/
utils.py
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
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