不同形状NumPy矩阵操作向量化实现类别概率提取的问题咨询

现有两个NumPy数组：

• 数组A（形状(10, 64, 10)）：维度含义为(类别数, 特征数, 分箱数)，代表10个类别、64个特征，每个特征的概率密度函数被划分为10个bin，最内层元素为对应概率值
• 数组B（形状(x, 64)）：共x个样本，每个元素对应第i个特征所属的bin索引

简化的数组A（形状(2,4,3)，对应2个类别、4个特征、3个分箱）：

[[[0.62, 0.,   0.  ],
  [0.12, 0.09, 0.01],
  [0.59, 0.01, 0.  ],
  [0.62, 0.,   0.  ]],
 [[0.62, 0.,   0.  ],
  [0.59, 0.01, 0.  ],
  [0.62, 0.,   0.  ],
  [0.62, 0.,   0.  ]]]

简化的数组B（形状(3,4)，对应3个样本、4个特征）：

[[0, 0, 2, 1],
 [0, 0, 1, 0],
 [0, 2, 1, 0]]

第一个样本的预期输出参考：

• 类别0：[0.62, 0.12, 0.00, 0.00]
• 类别1：[0.62, 0.59, 0.00, 0.00]

• 直接使用NumPy高级整数索引：你可以把数组B作为数组A最后一个维度（分箱维度）的索引，调整两个数组的维度适配广播规则，就可以一次性批量取出所有对应位置的概率值，不需要嵌套循环遍历样本和特征。
• 如果广播适配有问题，可以先对数组B做维度扩展，增加适配类别维度的轴之后再做索引取值。

• 没有利用NumPy的向量化特性，Python层级的嵌套for循环在样本量较大时性能损耗非常严重，效率远低于NumPy底层C实现的批量运算。
• 缺少索引合法性校验：没有判断数组B中的bin索引是否在数组A分箱数的合法范围内，遇到超出范围的索引会直接抛出异常，没有容错机制。
• 变量命名有歧义：bins变量名容易和NumPy本身的分箱相关功能混淆，建议改为更贴合语义的命名，比如sample_feature_bin_index。

内容的提问来源于stack exchange，提问作者Björn Bulkens