Shifting Elements in a Numpy Array

Question

I am trying to implement a function that will shift a particular element in a numpy 2d array a desired number of spaces in a given direction. Empty spaces should be filled with 0's. This function will take as input a numpy array, the x and y coordinates, a desired direction, and the number of spaces to shift.

For example, where shift is the hypothetical function:

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]]

I have found that I can achieve the desired shifting of all elements of either a row or column along that row or column with numpy's roll function. However, this approach simply cycles elements back to the beginning of the same row or column and does not fill empty spaces with 0's. For example:

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

Any assistance is very much appreciated.

edit: The x and y coordinates are the coordinates of the element to be shifted within the 2d array. The the rest of the elements within the same row or column are then shifted with respect to that element in the desired direction. For example shift(arr, 2, 2, 'down', 1) shifts the elements in the column with respect to the element at (2, 2) down by 1. All input values may be assumed to be valid at all times.

edit: This problem differs from the one linked in that elements are shifted with respect to the element at the coordinates provided, and this shifting occurs in a nested array. Furthermore, this solution does not allow for shifting of elements either up or down within the same column.

Answer 1

Just roll and then zero out data that got rotated around:

# Let direction = "down" and col = 0, n = 2
In [1]: i
Out[1]:
array([[1., 0., 0.],
       [1., 1., 1.],
       [0., 0., 1.]])

In [2]: i[:, col] = np.roll(i[:, col], n)

In [3]: i[:n, col] = 0

In [4]: i
Out[4]:
array([[0., 0., 0.],
       [0., 1., 1.],
       [1., 0., 1.]])

You'd need to implement the equivalent versions for the other three directions, but it'd just be variations of those two lines.

Answer 2

Here is a more or less comprehensive function to solve it:

def shift(a, i, j, dir, n, fill=0, inplace=False):
    out = a
    if not inplace:
        out = a.copy()
    if dir == 'down':
        if n < 0:
            return shift(out, i, j, 'up', -n, fill=fill, inplace=True)
        n = min(n, a.shape[0] - i)
        out[i+n:, j] = a[i:a.shape[0] - n, j]
        out[i:i+n, j] = fill
    elif dir == 'up':
        if n < 0:
            return shift(out, i, j, 'down', -n, fill=fill, inplace=True)
        n = min(n, i+1)
        out[:i+1-n, j] = a[n:i+1, j]
        out[i+1-n:i+1, j] = fill
    elif dir == 'right':
        if n < 0:
            return shift(out, i, j, 'left', -n, fill=fill, inplace=True)
        n = min(n, a.shape[1] - j)
        out[i, j+n:] = a[i, j:a.shape[1] - n]
        out[i, j:j+n] = fill
    elif dir == 'left':
        if n < 0:
            return shift(out, i, j, 'right', -n, fill=fill, inplace=True)
        n = min(n, j+1)
        out[i, :j+1-n] = a[i, n:j+1]
        out[i, j+1-n:j+1] = fill
    else:
        raise ValueError('Unknown direction "{}".'.format(dir))
    return out

Some tests:

import numpy as np

arr = np.arange(25).reshape((5, 5))
print(arr)
# [[ 0  1  2  3  4]
#  [ 5  6  7  8  9]
#  [10 11 12 13 14]
#  [15 16 17 18 19]
#  [20 21 22 23 24]]
print(shift(arr, 2, 1, 'up', 2))
# [[ 0 11  2  3  4]
#  [ 5  0  7  8  9]
#  [10  0 12 13 14]
#  [15 16 17 18 19]
#  [20 21 22 23 24]]
print(shift(arr, 2, 1, 'left', -2))
# [[ 0  1  2  3  4]
#  [ 5  6  7  8  9]
#  [10  0  0 11 12]
#  [15 16 17 18 19]
#  [20 21 22 23 24]]
print(shift(arr, 2, 1, 'down', 1, fill=100))
# [[  0   1   2   3   4]
#  [  5   6   7   8   9]
#  [ 10 100  12  13  14]
#  [ 15  11  17  18  19]
#  [ 20  16  22  23  24]]
shift(arr, 2, 1, 'right', 3, inplace=True)
print(arr)
# [[ 0  1  2  3  4]
#  [ 5  6  7  8  9]
#  [10  0  0  0 11]
#  [15 16 17 18 19]
#  [20 21 22 23 24]]

EDIT

Following discussion in comments, I add another function(s) to solve the problem of shifting "sliding tiles":

import numpy as np

def shift_vector(v, i, n, empty=0):
    if n < 0:
        return shift_vector(v[::-1], len(v) - i - 1, -n)[::-1]
    if n < len(v) - i:
        # Find n empty places after i
        idx = np.where(np.cumsum(v[i + 1:] == empty) == n)[0]
        last_zero_idx = idx[0] if len(idx) > 0 else len(v) - i - 1
        # Get non-empty values
        v_slice = v[i + 1:i + last_zero_idx + 1]
        values = v_slice[np.where(v_slice != empty)[0]]
        # Copy to vector
        v[i + n] = v[i]
        r = range(i + n + 1, min(i + last_zero_idx + 2, len(v)))
        v[r] = values[:len(r)]
    v[i:i + n] = empty
    return v

def shift(a, i, j, dir, n, empty=0, inplace=False):
    out = a
    if not inplace:
        out = a.copy()
    if dir == 'down':
        out[:, j] = shift_vector(out[:, j], i, n, empty=empty)
    elif dir == 'up':
        out[:, j] = shift_vector(out[:, j], i, -n, empty=empty)
    elif dir == 'right':
        out[i, :] = shift_vector(out[i, :], j, n, empty=empty)
    elif dir == 'left':
        out[i, :] = shift_vector(out[i, :], j, -n, empty=empty)
    else:
        raise ValueError('Unknown direction "{}".'.format(dir))
    return out


m = np.array([[1, 0, 0, 2],
              [3, 4, 0, 0],
              [5, 0, 6, 7],
              [0, 8, 9, 0]])
print("m")
print(m)
print("shift(m, 1, 0, 'right', 2)")
print(shift(m, 1, 0, 'right', 2))
print("shift(m, 3, 1, 'down', -2)")
print(shift(m, 3, 1, 'down', -2))
print("shift(m, 0, 3, 'left', 3)")
print(shift(m, 0, 3, 'left', 3))
print("shift(m, 2, 2, 'up', 1)")
print(shift(m, 2, 2, 'up', 1))

Output:

m
[[1 0 0 2]
 [3 4 0 0]
 [5 0 6 7]
 [0 8 9 0]]
shift(m, 1, 0, 'right', 2)
[[1 0 0 2]
 [0 0 3 4]
 [5 0 6 7]
 [0 8 9 0]]
shift(m, 3, 1, 'down', -2)
[[1 4 0 2]
 [3 8 0 0]
 [5 0 6 7]
 [0 0 9 0]]
shift(m, 0, 3, 'left', 3)
[[2 0 0 0]
 [3 4 0 0]
 [5 0 6 7]
 [0 8 9 0]]
shift(m, 2, 2, 'up', 1)
[[1 0 0 2]
 [3 4 6 0]
 [5 0 0 7]
 [0 8 9 0]]