Estou tentando implementar uma função que irá mudar um elemento específico em uma matriz 2d numpy um número desejado de espaços em uma determinada direção. Os espaços vazios devem ser preenchidos com zeros. Esta função terá como entrada uma matriz numpy, as coordenadas xey, uma direção desejada e o número de espaços a serem deslocados.

Por exemplo, onde shift é a função hipotética:

arr = np.array([1, 1, 1], [1, 1, 1], [1, 1, 1])
arr out: [[1, 1, 1],
          [1, 1, 1],
          [1, 1, 0]]

shift(arr, 0, 0, 'right', 2)
arr out: [[0, 0, 1],
          [1, 1, 1],
          [1, 1, 0]]

shift(arr, 0, 2, 'down', 1)
arr out: [[0, 0, 0],
          [1, 1, 1],
          [1, 1, 1]]

Descobri que posso obter a mudança desejada de todos os elementos de uma linha ou coluna ao longo dessa linha ou coluna com a função roll de numpy. No entanto, essa abordagem simplesmente alterna elementos de volta ao início da mesma linha ou coluna e não preenche espaços vazios com zeros. Por exemplo:

arr[:, 0] = np.roll(arr[:, 0], 1)
arr out: [[1, 0, 0],
          [0, 1, 1],
          [1, 1, 1]]

Qualquer assistência é muito apreciada.

editar: As coordenadas x e y são as coordenadas do elemento a ser deslocado dentro da matriz 2d. O restante dos elementos na mesma linha ou coluna é então deslocado em relação a esse elemento na direção desejada. Por exemplo, shift (arr, 2, 2, 'down', 1) desloca os elementos na coluna em relação ao elemento em (2, 2) para baixo em 1. Todos os valores de entrada podem ser considerados válidos o tempo todo.

editar: esse problema difere daquele vinculado em que os elementos são deslocados em relação ao elemento nas coordenadas fornecidas, e essa mudança ocorre em uma matriz aninhada. Além disso, esta solução não permite a troca de elementos para cima ou para baixo na mesma coluna.