python - 如何将 numpy 的整数数组索引应用于 ndarray-s 和 python 列表？

Question

Numpy 有 "integer array indexig" ，它允许使用其他 ndarray 或列表索引 ndarray:

>> A = np.arange(10,20)
>> A[[1,2,3]]
Out[14]: array([11, 12, 13])

假设我不知道A是什么，它可以是ndarray或Python的列表。

numpy 中是否有任何显式索引函数，它允许相同的索引并接受两种类型？

例如:

>> A = np.arange(10,20)
>> np.get_elements(A, [1,2,3])
Out[14]: array([11, 12, 13])
>> A = range(10,20)
>> np.get_elements(A, [1,2,3])
Out[15]: [11, 12, 13]

Answer 1

def get_elements(A, idx):
    try:
        return A[idx]
    except TypeError:
        import operator
        return list(operator.itemgetter(*idx)(A))

In [35]: get_elements(np.arange(10), [1,3,4])
Out[35]: array([1, 3, 4])
In [36]: get_elements(np.arange(10).tolist(), [1,3,4])
Out[36]: [1, 3, 4]

itemgetter 只是一个便利类。列表理解也同样好

In [39]: [A[i] for i in [1,3,4]]
Out[39]: [1, 3, 4]

像 take 这样的 numpy 函数返回一个数组

In [40]: np.take(A,[1,3,4])
Out[40]: array([1, 3, 4])

如果需要，可以将其转换回列表。