Выберите из трехмерного массива двумерный массив
У меня есть два массива:
a: трехмерный исходный массив (Н х М х 2)b: массив двумерных индексов (Н х М), содержащий 0 и 1.
Я хочу использовать индексы в b для выбора соответствующих элементов a в его третьем измерении. Результирующий массив должен иметь размеры Н х М. Вот пример в виде кода:
import numpy as np
a = np.array( # dims: 3x3x2
[[[ 0, 1],
[ 2, 3],
[ 4, 5]],
[[ 6, 7],
[ 8, 9],
[10, 11]],
[[12, 13],
[14, 15],
[16, 17]]]
)
b = np.array( # dims: 3x3
[[1, 1, 1],
[1, 1, 1],
[1, 1, 1]]
)
# select the elements in a according to b
# to achieve this result:
desired = np.array(
[[ 1, 3, 5],
[ 7, 9, 11],
[13, 15, 17]]
)
Сначала я думал, что у этого должно быть простое решение, но я вообще не мог его найти. Поскольку я хотел бы перенести его на тензорный поток, я был бы признателен, если кто-нибудь знает для этого решение типа numpy.
Редактировать: Третье измерение a может содержать более двух элементов. Следовательно, b также может содержать индексы, отличные от 0 и 1 — это не логическая маска.
@ParitoshSingh Для каждого элемента x в позиции (m, n) в b выходной массив в позиции (m, n) должен быть равен a[m, n, x].
Ответы 4
Как предлагает @jdehesa, мы можем использовать np.ogrid для получения индексов для первых двух осей:
ax0, ax1 = np.ogrid[:b.shape[0], :b.shape[1]]
И затем мы можем использовать b для прямого индексирования по последней оси. Обратите внимание, что ax0 и ax1 будут транслироваться в форму b:
desired = a[ax0, ax1 ,b]
print(desired)
array([[ 1, 3, 5],
[ 7, 9, 11],
[13, 15, 17]])
Вы также можете получить то же самое с ax0, ax1 = np.ogrid[:b.shape[0], :b.shape[1]], а затем desired = a[ax0, ax1, b] (трансляция индексов).
Ах да, так намного лучше, спасибо за предложение @jdehesa
Мы можем использовать np.where для этого:
np.where(b, a[:, :, 1], a[:, :, 0])
Выход:
array([[ 1, 3, 5],
[ 7, 9, 11],
[13, 15, 17]])
Мне нравится, как вы используете здесь np.where, но последнее измерение a может содержать более двух элементов, поэтому мне действительно нужны числовые индексы из b не только как логическая маска. Извините, я должен уточнить это в своем вопросе.
@mrzo Тогда ответ take_along_axis, скорее всего, то, что вы хотите.
Вы можете использовать np.take_along_axis:
import numpy as np
a = np.array(
[[[ 0, 1],
[ 2, 3],
[ 4, 5]],
[[ 6, 7],
[ 8, 9],
[10, 11]],
[[12, 13],
[14, 15],
[16, 17]]])
b = np.array(
[[1, 1, 1],
[1, 1, 1],
[1, 1, 1]])
print(np.take_along_axis(a, b[..., np.newaxis], axis=-1)[..., 0])
# [[ 1 3 5]
# [ 7 9 11]
# [13 15 17]]
Я добавил несколько решений для tensorflow.
import tensorflow as tf
a = tf.constant([[[ 0, 1],[ 2, 3],[ 4, 5]],
[[ 6, 7],[ 8, 9],[10, 11]],
[[12, 13],[14, 15],[16, 17]]],dtype=tf.float32)
b = tf.constant([[1, 1, 1],[1, 1, 1],[1, 1, 1]],dtype=tf.int32)
# 1. use tf.gather_nd
colum,row = tf.meshgrid(tf.range(a.shape[0]),tf.range(a.shape[1]))
idx = tf.stack([row, colum, b], axis=-1) # Thanks for @jdehesa's suggestion
result1 = tf.gather_nd(a,idx)
# 2. use tf.reduce_sum
mask = tf.one_hot(b,depth=a.shape[-1],dtype=tf.float32)
result2 = tf.reduce_sum(a*mask,axis=-1)
# 3. use tf.boolean_mask
mask = tf.one_hot(b,depth=a.shape[-1],dtype=tf.float32)
result3 = tf.reshape(tf.boolean_mask(a,mask),b.shape)
with tf.Session() as sess:
print('method 1: \n',sess.run(result1))
print('method 2: \n',sess.run(result2))
print('method 3: \n',sess.run(result3))
method 1:
[[ 1. 3. 5.]
[ 7. 9. 11.]
[13. 15. 17.]]
method 2:
[[ 1. 3. 5.]
[ 7. 9. 11.]
[13. 15. 17.]]
method 3:
[[ 1. 3. 5.]
[ 7. 9. 11.]
[13. 15. 17.]]
Хороший. Просто в качестве предложения вы можете сделать idx = tf.concat([row[:,:,tf.newaxis],colum[:,:,tf.newaxis],b[:,:,tf.newaxis]],axis=-1) немного более кратким с idx = tf.stack([row, colum, b], axis=-1).
не могли бы вы уточнить логику выбора? я не вижу никакой связи между
bи конечным результатом.