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具有相同的索引（索引= 3），與A中的元素13和B中的元素8相同。

Answer 1

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

返回與元素具有相同形狀的布爾數組，如果元素的元素在test_elements中，則返回True，否則返回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是否轉換為集合，但是我想它將比您當前的代碼更快。

Answer 2

我不確定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)