python - numpy.choose 的替代方案允许任意参数或至少超过 32 个参数？

Question

在我的代码中，我遇到了 numpy.choose 方法不接受所有参数的问题，因为它受到 NPY_MAXARGS 的限制(= 32)。是否有可用的替代方案，允许任意数量的参数数组或至少超过 32 并且与 numpy.choose 一样快？

choices = [np.arange(0,100)]*100
selection = [0] * 100
np.choose(selection, choices)

>> ValueError: Need between 2 and (32) array objects (inclusive).

任何帮助将不胜感激... :)

Answer 1

索引可以作为列表给出。假设 selections 与 choices 的长度相同:

b = numpy.array(choices)
result = b[range(len(selections)), selections]

将给出由选择中的索引指定的选择中的值。查看实际效果:

numpy.random.seed(1)
b = numpy.random.randint(0,100,(5,10))
>>> array([[37, 12, 72,  9, 75,  5, 79, 64, 16,  1],
          [76, 71,  6, 25, 50, 20, 18, 84, 11, 28],
          [29, 14, 50, 68, 87, 87, 94, 96, 86, 13],
          [ 9,  7, 63, 61, 22, 57,  1,  0, 60, 81],
          [ 8, 88, 13, 47, 72, 30, 71,  3, 70, 21]])

selections = numpy.random.randint(0,10,5)
>>> array([1, 9, 3, 4, 8])

result = b[range(len(selections)),selections]
>>>> array([12, 28, 68, 22, 70])