英文:
Slice ndarray with indexes defined in another ndarray
问题
有没有一种方法可以使用来自idx的第i行作为a的第i行的索引，即a[i,idx[i]]？
可以通过循环轻松完成这个操作：
np.vstack([a[i, idx[i]] for i in range(len(a))])
但我认为肯定有一种使用NumPy的方法来实现这个目标。
最终结果应该是：
[[1, 3],
[5, 6]]
使用a[idx]不起作用。如果我使用a[:,idx]，那么我会得到太多的结果，即：
array([[[1, 3],
[2, 3]],
[[4, 6],
[5, 6]]])
这接近但不正确。
英文:
Say I have the following:
a = np.array([[1,2,3],[4,5,6]])
idx = np.array([[0,2],[1,2]])
Is there a way to "numpy-slice" the array such that I use the i-th row from idx as index for the i-th row from a i.e a[i,idx[i]]?
This can easily be done by a loop:
np.vstack([a[i, idx[i]] for i in range(len(a))])
But I thought there must be a numpy way of doing it.
The end-result should be:
[[1,3],
[5,6]]
Doing a[idx] does not work. If I do a[:,idx] then I get too many results i.e that gives
array([[[1, 3],
[2, 3]],
[[4, 6],
[5, 6]]])
which is close, but not correct.
答案1
得分: 2
使用 np.take_along_axis 在 1 轴上选择所需的索引：
>>> np.take_along_axis(a, idx, 1)
array([[1, 3],
[5, 6]])
英文:
Use np.take_along_axis to select the required indices along the 1 axis:
>>> np.take_along_axis(a, idx, 1)
array([[1, 3],
[5, 6]])
答案2
得分: 0
你可以使用NumPy的高级索引来获得你想要的结果。
import numpy as np
a = np.array([[1,2,3],[4,5,6]])
idx = np.array([[0,2],[1,2]])
res = a[np.arange(a.shape[0])[:,None], idx]
print(res)
输出：
[[1 3]
[5 6]]
英文:
You can use NumPy's advanced indexing to get your desired result.
import numpy as np
a = np.array([[1,2,3],[4,5,6]])
idx = np.array([[0,2],[1,2]])
res = a[np.arange(a.shape[0])[:,None], idx]
print(res)
Output:
[[1 3]
[5 6]]
通过集体智慧和协作来改善编程学习和解决问题的方式。致力于成为全球开发者共同参与的知识库，让每个人都能够通过互相帮助和分享经验来进步。
评论