Python - 优化代码(遍历数组)

Question

我有两个数组，A 和 B。基于一个函数，我从 A 创建一个子集 C，从 B 创建一个子集 D。现在，我想将它们连接在一起(我想保留 C 中的那些项目和 D 在 A 和 B 中具有相同的索引)。

A = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20])
B = np.array([20,19,18,17,16,15,14,13,12,11,10,9,8,7,6,5,4,3,2,1])
C = np.array([3,7,13])
D = np.array([18,8])

fvl1 = list()
fvl2 = list()


for i in C:
    index = np.where(A == i)
    if B[index] in D:
        fvl1.append(A[index])
        fvl2.append(B[index])

print(fvl1)
print(fvl2)

Output:
[array([3]), array([13])]
[array([18]), array([8])]

这是我到目前为止所做的，但它真的很慢。有什么方法可以优化它以使其运行得更快吗？

LE: 你可以看到 A 中的元素 3 与 B 中的元素 18 具有相同的索引(index = 3)，与 A 中的元素 13 和 B 中的元素 8 相同。

Answer 1

你可以使用 np.ind1d或 np.isin :

Returns a boolean array of the same shape as element that is True where an element of element is in test_elements and False otherwise.

>>> import numpy as np
>>> A = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20])
>>> B = np.array([20,19,18,17,16,15,14,13,12,11,10,9,8,7,6,5,4,3,2,1])
>>> C = np.array([3,7,13])
>>> D = np.array([18,8])
>>> np.in1d(A, C) & np.in1d(B, D)
array([False, False,  True, False, False, False, False, False, False,
       False, False, False,  True, False, False, False, False, False,
       False, False], dtype=bool)
>>> indices = np.isin(A, C) & np.isin(B, D)
>>> A[indices]
array([ 3, 13])
>>> B[indices]
array([18,  8])

尚不清楚 C 和 D 是否转换为集合，但我想它会比您当前的代码更快。