Произвольно перемещать выбранные элементы массива numpy в соответствии с распределением? (Расширить размеры 2D-массива до 3D случайными числами)

Question

У меня есть массив 2D numpy, состоящий из 0 и 1.

a = np.array([[1, 0, 0, 1], [0, 1, 0, 1]])

Мне нужно создать новый массив a_new в соответствии с этим:
Для каждого 1 в местоположении [l, k] массива a выберите случайное число в соответствии с желаемым распределение (например, shift = np.int64(np.ceil(np.random.gamma(1, 3)))) и установите его на a_new[l, k, shift], если смещение меньше, чем N, в противном случае игнорируйте это 1.

Вот его имплантация oop. Есть ли более быстрое (возможно, операция с массивом) решение этой проблемы? Размер матрицы a большой.

import numpy as np

N = 5
a = np.array([[1, 0, 0, 1], [0, 1, 0, 1]])
a_new = np.zeros((a.shape[0], a.shape[1], N))
for k in np.arange(a.shape[1]):
    for l in np.arange(a.shape[0]):
        if a[l, k]:
            shift = np.int64(np.ceil(np.random.gamma(1, 3)))
            if (shift < N):
                a_new[l, k, shift] = 1

образец выборки:

a 
[[1 0 0 1]
 [0 1 0 1]]
a_new 
[[[0. 0. 1. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 1. 0. 0.]]

 [[0. 0. 0. 0. 0.]
  [0. 0. 0. 1. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]]

Answer 1

Вот трюк, использующий np.bincount:

a = np.array([[1, 0, 0, 1], [0, 1, 0, 1]])
N = 5
X,Y = a.nonzero()
Z = np.ceil(np.random.gamma(1,3,X.shape)).astype(int)
Z
# array([ 2,  1, 15,  2])
flat = np.ravel_multi_index((X[Z<N],Y[Z<N],Z[Z<N]),a.shape+(N,))
np.bincount(flat,None,a.size*N).reshape(*a.shape,N)
# array([[[0, 0, 1, 0, 0],
#         [0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0]],
#
#        [[0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0],
#         [0, 0, 1, 0, 0]]])

ОБНОВЛЕНИЕ: С кратностями:

a = np.array([[1, 0, 0, 3], [0, 10, 0, 1]])

N = 5
X,Y = a.nonzero()
times = a[X,Y]
X = X.repeat(times)
Y = Y.repeat(times)
Z = np.ceil(np.random.gamma(1,3,X.shape)).astype(int)
flat = np.ravel_multi_index((X[Z<N],Y[Z<N],Z[Z<N]),a.shape+(N,))
np.bincount(flat,None,a.size*N).reshape(*a.shape,N)
# array([[[0, 1, 0, 0, 0],
#        [0, 0, 0, 0, 0],
#        [0, 0, 0, 0, 0],
#        [0, 1, 0, 0, 0]],
#
#       [[0, 0, 0, 0, 0],
#        [0, 4, 2, 0, 3],
#        [0, 0, 0, 0, 0],
#        [0, 0, 0, 0, 0]]])

Answer 2

Вы также можете попробовать следующее, что может помочь,

import numpy as np
N = 5
a = np.array([[1, 0, 0, 1], [0, 1, 0, 1]])
row, col = np.where(a==1)
a_new = np.zeros([a.shape[0], a.shape[1], N])
a_new[row,col] = (np.int64(np.ceil(np.random.gamma(1,3,[a.shape[1],N])))<N)

Answer 3

Я нашел решение с расширенной индексацией в numpy, которое значительно быстрее. По сути, мы можем использовать ненулевые элементы a в качестве индексов для первых двух измерений a_new и случайные значения shift в качестве индексов для третьего измерения a_new с некоторой предварительной фильтрацией (удалить из привязанного случайного числа) числа, и преобразование массива shift в форму, отличную от ненулевого подмассива a. Вот рабочий код:
import numpy как np

N = 5
a = np.array([[1, 0, 0, 1], [0, 1, 0, 1]])
a_new = np.zeros((a.shape[0], a.shape[1], N))
shift = np.int64(np.ceil(np.random.gamma(1, 3, a.shape)))
a[shift > N-1] = 0
shift[shift > N-1] = 0
shift = (shift * a).reshape(1, -1)
shift = shift[shift > 0]
non_zero_a = np.nonzero(a)
a_new[non_zero_a[0], non_zero_a[1], shift] = 1