python - 使用广播过滤带有索引的多维数组

Question

我想知道是否有一种方法可以使用索引列表过滤多维数组。

例如，如果有一个多维数组，如[[40,60,15],[25,30,45]]，我想获得使用以下索引[(0,2),(1,0)]进行过滤时的结果值。请注意，索引是一个元组列表

注意:我正在尝试找出最快的方法。我设法使用 for 循环和 map + lambda 函数来做到这一点， 但是想知道 numpy 广播是否可行？

 from numpy import np

 some_arr = np.array([40,60,15],[25,30,45])
 some_idx = [(0,2),(1,0)]
 
 #Method 1 : Using for loop
 new_list = []
 for idx in some_idx:
     new_list.append(some_arr[idx])

 # Method 2: Using lambda
 new_list = list(map(lambda idx : some_arr[idx], some_idx))

 # Method 3 : Using numpy ??? (i get an error ofcourse but this is what I tried to do)
 new_arr = some_arr[some_idx]

我最终遇到了方法 3 的错误:索引 2 超出了尺寸为 2 的轴 0 的范围 此示例的输出应为 [15,25]

Answer 1

使用 zip 代替 ((x_0, x_1...),(y_0,y_1...)) 索引元组，然后您可以使用它沿两个轴执行整数数组索引:

some_arr = np.array([[40,60,15],[25,30,45]])
some_idx = [(0,2),(1,0)]

some_arr[tuple(zip(*some_idx))]
# array([15, 25])

---

与第一种方法比较:

new_list = []
for idx in some_idx:
    new_list.append(some_arr[idx])

np.array_equal(new_list, some_arr[tuple(zip(*some_idx))])
# True