Роллинг строк матрицы независимо

У меня есть матрица (2d numpy ndarray, если быть точным):

A = np.array([[4, 0, 0],
              [1, 2, 3],
              [0, 0, 5]])

И я хочу катить каждую строку A независимо, в соответствии с значениями roll в другом массиве:

r = np.array([2, 0, -1])

То есть, я хочу сделать это:

print np.array([np.roll(row, x) for row,x in zip(A, r)])

[[0 0 4]
 [1 2 3]
 [0 5 0]]

Есть ли способ сделать это эффективно? Возможно, используя причудливые индексирующие трюки?

ОТВЕТЫ

Ответ 1

Конечно, вы можете сделать это с помощью расширенного индексирования, возможно, это самый быстрый способ, вероятно, зависит от размера вашего массива (если ваши строки большие, это может быть не так):

rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]

# Use always a negative shift, so that column_indices are valid.
# (could also use module operation)
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:, np.newaxis]

result = A[rows, column_indices]

Ответ 2

numpy.lib.stride_tricks.as_strided stricks (аббревиатура, предназначенная для каламбура) снова!

Говоря о причудливых индексационных трюках, там печально известный - np.lib.stride_tricks.as_strided. Идея/трюк заключалась бы в том, чтобы получить отрезанную часть, начиная с первого столбца, до второй последней и конкатенации в конце. Это гарантирует, что мы можем двигаться вперед, когда это необходимо, чтобы использовать np.lib.stride_tricks.as_strided и, таким образом, избежать необходимости фактического откат. Это вся идея!

Теперь, с точки зрения фактической реализации, мы будем использовать scikit-image view_as_windows для элегантного использования np.lib.stride_tricks.as_strided под капюшонами. Таким образом, окончательная реализация будет -

from skimage.util.shape import view_as_windows as viewW

def strided_indexing_roll(a, r):
    # Concatenate with sliced to cover all rolls
    a_ext = np.concatenate((a,a[:,:-1]),axis=1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = a.shape[1]
    return viewW(a_ext,(1,n))[np.arange(len(r)), (n-r)%n,0]

Здесь образец пробега -

In [327]: A = np.array([[4, 0, 0],
     ...:               [1, 2, 3],
     ...:               [0, 0, 5]])

In [328]: r = np.array([2, 0, -1])

In [329]: strided_indexing_roll(A, r)
Out[329]: 
array([[0, 0, 4],
       [1, 2, 3],
       [0, 5, 0]])

Бенчмаркинг

# @seberg solution
def advindexing_roll(A, r):
    rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]    
    r[r < 0] += A.shape[1]
    column_indices = column_indices - r[:,np.newaxis]
    return A[rows, column_indices]

Позвольте сделать некоторый бенчмаркинг по массиву с большим количеством строк и столбцов -

In [324]: np.random.seed(0)
     ...: a = np.random.rand(10000,1000)
     ...: r = np.random.randint(-1000,1000,(10000))

# @seberg solution
In [325]: %timeit advindexing_roll(a, r)
10 loops, best of 3: 71.3 ms per loop

#  Solution from this post
In [326]: %timeit strided_indexing_roll(a, r)
10 loops, best of 3: 44 ms per loop

Ответ 3

Если вам нужно более общее решение (с любой формой и любой осью), я изменил решение @seberg:

def indep_roll(arr, shifts, axis=1):
    """Apply an independent roll for each dimensions of a single axis.

    Parameters
    ----------
    arr : np.ndarray
        Array of any shape.

    shifts : np.ndarray
        How many shifting to use for each dimension. Shape: '(arr.shape[axis],)'.

    axis : int
        Axis along which elements are shifted. 
    """
    arr = np.swapaxes(arr,axis,-1)
    all_idcs = np.ogrid[[slice(0,n) for n in arr.shape]]

    # Convert to a positive shift
    shifts[shifts < 0] += arr.shape[-1] 
    all_idcs[-1] = all_idcs[-1] - shifts[:, np.newaxis]

    result = arr[tuple(all_idcs)]
    arr = np.swapaxes(result,-1,axis)
    return arr