python - 在 numpy 中从多个集合中选取并分配多个子集

Question

给定形状为 (2, 3, 4) 的基本数组 X，它可以解释为两组，每组 3 个元素，其中每个元素都是 4 维，我想通过以下方式从此数组 X 中进行采样。我想从 2 个集合中的每一个中挑选 2 个子集，每个子​​集由长度为 3 的二进制数组定义，其他子集将设置为 0。因此采样过程由形状 (2, 2, 3) 的数组定义。此采样的结果应具有形状 (2, 2, 3, 4)。

这是执行我需要的代码，但我想知道它是否可以使用 numpy 索引更有效地重写。

import numpy as np
np.random.seed(3)

sets = np.random.randint(0, 10, [2, 3, 4])
subset_masks = np.random.randint(0, 2, [2, 2, 3])

print('Base set\n', sets, '\n')
print('Subset masks\n', subset_masks, '\n')

result = np.empty([2, 2, 3, 4])
for set_index in range(sets.shape[0]):
    for subset_index, subset in enumerate(subset_masks[set_index]):
        print('----')
        picked_subset = subset.reshape(3, 1) * sets[set_index]
        result[set_index][subset_index] = picked_subset
        print('Picking subset ', subset, 'from set #', set_index)
        print(picked_subset, '\n')

输出

Base set
 [[[8 9 3 8]
  [8 0 5 3]
  [9 9 5 7]]

 [[6 0 4 7]
  [8 1 6 2]
  [2 1 3 5]]] 

Subset masks
 [[[0 0 1]
  [1 0 0]]

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

----
Picking subset  [0 0 1] from set # 0
[[0 0 0 0]
 [0 0 0 0]
 [9 9 5 7]] 

----
Picking subset  [1 0 0] from set # 0
[[8 9 3 8]
 [0 0 0 0]
 [0 0 0 0]] 

----
Picking subset  [1 0 1] from set # 1
[[6 0 4 7]
 [0 0 0 0]
 [2 1 3 5]] 

----
Picking subset  [0 1 1] from set # 1
[[0 0 0 0]
 [8 1 6 2]
 [2 1 3 5]] 

Answer 1

通过沿着最后一个轴添加 subset_masks 的新轴并将 sets 添加为第二个轴，将每个轴扩展到 4D。为了添加这些新轴，我们可以使用 None/np.newaxis。然后，杠杆NumPy broadcasting执行逐元素乘法，如下所示 -

subset_masks[...,None]*sets[:,None]

也许只是为了好玩，我们也可以使用 np.einsum -

np.einsum('ijk,ilj->iljk',sets,subset_masks)