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