Как вернуть список элементов из массива заданного списка по индексам в 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)]