如何从给定的numpy索引列表中的数组中返回项目列表

Question

我有一个数组arr和一个我想获取的索引列表indices。我想获取与indices中的项目相对应的数组子集，以及该数组的补数。

例如

对于

arr = np.asarray([2, 4, 1, 5, 6])
indices = np.asarray([2, 4])

我会得到

[1, 6] and [2, 4, 5]

谢谢

Answer 1

使用np.isin或np.in1d（使用掩码）：

arr = np.asarray([2, 4, 1, 5, 6])
indices = np.asarray([2, 4])
m = np.in1d(np.arange(len(arr)), indices)
arr1, arr2 = arr[m], arr[~m]
arr1, arr2
>>array([1, 6]), array([2, 4, 5])

或者，对补充部分使用np.setdiff1d（对于较大的数组和索引，可以更快）：

arr1 = arr[indices]
arr2 = arr[np.setdiff1d(np.arange(len(arr)), indices, True)]