Question

我有一个数组，让我们说arr = [1,2,3,4,5,6,7,8] 和另一个索引数组：idx = [0,3,4,6]

我希望得到两个数组，一个只是来自arr的那些索引：[1,4,5,7] 剩下的就是另一个：[2,3,6,8]

有人可以帮我吗？我只能想到做丑陋的方法，但它必须是一些优雅的功能。

非常感谢！

Answer 1

你可以这样做：

selected = [arr[i] for i in idx]
other = [v for i, v in enumerate(arr) if i not in idx]

如果arr没有重复项，您也可以这样做：

other = [v for v in arr if v not in selected]

Answer 2

方法：

a1 = [arr[x] for x in idx] 
a2 = [x for x in arr if x not in a1]

Answer 3

一次遍历：

no, yes = both = [], []
for i, x in enumerate(arr):
    both[i in idx].append(x)

或（由Chris_Rands评论）：

yes, no = [], []
for i, x in enumerate(arr):
    (yes if i in idx else no).append(x)

虽然idx应该是小的，或者变成一个集合（其他答案中的解决方案也是如此，我想他们只是不想谈论它。）

演示：

>>> if 1:
    arr = [1, 2, 3, 4, 5, 6, 7, 8]
    idx = [0, 3, 4, 6]
    no, yes = both = [], []
    for i, x in enumerate(arr):
        both[i in idx].append(x)
    print('yes', yes)
    print('no', no)

yes [1, 4, 5, 7]
no [2, 3, 6, 8]

Answer 4

numpy有一个简洁的解决方案：

import numpy as np

arr = np.asarray([1, 2, 3, 4, 5, 6, 7, 8])  # converts your list in numpy array
idx1 = [0, 3, 4, 6]
idx2 = [1, 2, 5, 7]

arr1 = arr[idx1]  # [1 4 5 7]
arr2 = arr[idx2]  # [2 3 6 8]

Answer 5

您可以将itertools用于单行解决方案：

import itertools
arr = [1, 2, 3, 4, 5, 6, 7, 8]
idx = [0, 3, 4, 6]
[(out_index, not_in_arr), (in_index, in_arr)] = [(a, list(b)) for a, b in itertools.groupby(sorted(arr, key=lambda x:arr.index(x) in idx), key=lambda x:arr.index(x) in idx)]
print(not_in_arr)
print(in_arr)

输出：

[2, 3, 6, 8]
[1, 4, 5, 7]

Answer 6

您还可以将arr中的每个值映射到字典，指示idx中是否存在该索引：

arr = [1, 2, 3, 4, 5, 6, 7, 8]
idx = [0, 3, 4, 6]

# converted to a set
idx_lookup = set(idx)

d = {x: True if i in idx_lookup else False for i, x in enumerate(arr)}
print(d)

这给出了这个词典：

{1: True, 2: False, 3: False, 4: True, 5: True, 6: False, 7: True, 8: False}

我还将idx转换为集合，因为在这种情况下，重复索引不是必需的，而set / dictionary lookup是O(1)。但是，列表查找是O(n)，因此如果可能，这种优化是值得的。

一旦你有这本词典，你可以filter出你想要保留的元素，以及其他元素：

keep = list(filter(lambda x: d[x], arr))
rest = list(filter(lambda x: not d[x], arr))

print(keep)
print(rest)

哪个输出：

[1, 4, 5, 7]
[2, 3, 6, 8]

注意：您还可以使用列表推导来过滤keep和rest：

keep = [x for x in arr if d[x]]
rest = [x for x in arr if not d[x]]

Python使用索引数组将数组切割成两个数组

6 个答案: