按列索引从NumPy数组中选择每行的一个元素

Question

有没有更好的方法从"input_array"和"select_id"中获取"output_array"?

我们能否摆脱range(input_array.shape[0])?

>>> input_array = numpy.array( [ [3,14], [12, 5], [75, 50] ] )
>>> select_id = [0, 1, 1]
>>> print input_array
[[ 3 14]
 [12  5]
 [75 50]]

>>> output_array = input_array[  range( input_array.shape[0] ), select_id ]
>>> print output_array
[ 3  5 50]

- Bystander

---

这种方法并不好，绝对不如你现有的方法好，但是 np.diagonal(input_array[:, select_id]) 也可以得到 array([ 3, 5, 50])。 - Jaime

除了使用 arange 替代 range 之外，问题中的高级索引解决方案已经是最佳选项了。 - user2357112

Answer 1

您可以使用numpy.choose从一个索引数组（在此案例中为select_id）和一组数组（在此案例中为input_array）构造数组以进行选择。但是，您可能需要先转置input_array以匹配维度。以下是一个小例子：

In [101]: input_array
Out[101]: 
array([[ 3, 14],
       [12,  5],
       [75, 50]])

In [102]: input_array.shape
Out[102]: (3, 2)

In [103]: select_id
Out[103]: [0, 1, 1]

In [104]: output_array = np.choose(select_id, input_array.T)

In [105]: output_array
Out[105]: array([ 3,  5, 50])

- mg007

---

我们如何就地修改这些值，而不是将它们输出？ - syllogismos

你可以使用这个链接：https://dev59.com/Bmsz5IYBdhLWcg3wn5fa - Steven

必须评论一下，这对于大型数组是不起作用的，正如Nathan在下面指出的那样。它会出现“ValueError：至少需要1个且最多32个数组对象。” 有人知道为什么这种方法局限于小型数组吗？ - Tony

Answer 2

请注意，numpy.choose仅在您有32个或更少的选择时才有效（在这种情况下，沿着您进行索引的数组的维度必须小于等于32）。此外，关于numpy.choose的文档中指出：

为了减少误解的机会，即使以下“滥用”名义上得到支持，也不应将choices视为单个数组，最外层类似序列的容器应该是列表或元组。

OP问道：

1. 是否有更好的方法从input_array和select_id获取output_array？
   • 我认为你最初提出的方式是这里介绍的最佳方法。它易于理解，可扩展到大型数组，并且高效。
2. 我们能否摆脱range(input_array.shape[0])？
   • 是的，如其他答案所示，但被接受的答案在一般情况下效果不佳，而OP已经建议采取的方法是好的。

- Nathan

---

那将是一个非常长的注释 ;) - Lincoln Bergeson

哈哈，当一个人写下一些内容作为答案时，我觉得它只会变得更长。然后我感到有责任通过直接回答原帖的问题来使它成为“答案”。 - Nathan

Answer 3

我认为枚举函数很方便。

[input_array[enum, item] for enum, item in enumerate(select_id)]

- y4suyuki

---

利用内联循环节约时间总是很好的。虽然我真的需要NumPy来处理大量数据... - Bystander

Answer 4

怎么样？

[input_array[x,y] for x,y in zip(range(len(input_array[:,0])),select_id)]

- Lee

---

我需要处理大量的数据，因此我尝试使用NumPy进行向量化计算。 - Bystander