Блог

У меня есть 3-мерный массив numpy с размерами (x = 2, y = 2, z = 3), как показано ниже

 a = [[[0,1,2],[3,4,5]],
    [[6,7,8],[9,10,11]]]
 

Я хочу получить первые N элементов из каждого (x, y) элемента a, где N определено в другом массиве размером x, y. Например.

 b = [[1,2],
     [0,0]]
 

Результат должен быть

 c = [[[0,1],[3,4,5]],
    [[6],[9]]
 

Как я могу сделать это без циклов?

Вы можете создать маску над элементами, a которые находятся в пределах заданных диапазонов, и использовать эту маску всякий раз, когда вам нужно работать a . Если вы хотите инкапсулировать маску и массив в одном объекте, взгляните на numpy.ma модуль.

 a = np.array([[[0,1,2],[3,4,5]],
             [[6,7,8],[9,10,11]]])

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

mask = np.arange(3)[None,None,:] <= b[:,:,None]

a[mask]
 

Вывод:

 array([0, 1, 3, 4, 5, 6, 9])
 

Но если вы хотите получить результат в виде неоднородного массива, я не могу придумать лучшего способа сделать это, чем использовать циклы.