Python-优化代码(通过数组循环)

Question

我有两个阵列，A和B。根据函数，我从A和一个子集D创建一个子集C和A子集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中的元素相同的索引（索引= 3），与a和a的元素13相同，来自b。

您可以使用np.ind1d或np.isin：

返回与元素相同形状的布尔数组 elements中的元素元素在test_lements中，否则为false。

>>> 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是否转换为集合，但我想它会比您当前的代码更快。

我不是100％确定我正确解释了您的问题，但是zip和列表理解的组合可能是您追求的。

fv = [(a, b) for (a, b) in zip(A, B)
      if <condition for first list>
      and <condition for second list>]
fv1 = [t[0] for t in fv]
fv2 = [t[1] for t in fv]

zip将您的列表加入到一个元组列表中。

list(zip([1, 2, 3],['a','b','c'])) == [(1,'a'),(2,'b'),(3,'c')]

您可以通过列表理解中的条件过滤所需的输出。例如：

fv = [(a, b) for (a, b) in zip(range(10),'abcdefghijklmnop') 
      if a%2 
      and b>'d']
[(5, 'f'), (7, 'h'), (9, 'j')]

您的情况实际上可以用作一个衬里：

fv1, fv2 = zip(*((a, b) for (a, b) in zip(A, B) if a in C and b in D))
fv1 #(3, 13)
fv2 #(18, 8)