使用列表推导式代替for循环中的append方法。

Question

我想使用列表推导来替代append方法，因为我知道这能使代码运行得更快。但当我面对numpy数组时，这个方法似乎不起作用。

import numpy as np

a = np.array([1, 2 , 4])
b = np.array([6, 5 , 2])
A = []
B = []
for i in range(0, 3):
    if a[i] > value1 and b[i] > value1:
        A.append(a[i])
        B.append(b[i])

例如，我尝试了这样的写法：

A = [a[i] for i in range(1, 3) if a[i] > value1 and b[i] > value1]

最后，我想知道是否有一种方式可以定义一个通用函数，根据类似上述示例中的条件，从一组数字的数组中挑选元素并添加到新列表中。

melpomene · Answer

numpy数组没有append方法。同时，使用for循环遍历numpy数组也是低效的。相反，你应该利用向量化操作和广播机制。

条件判断可以通过以下方式实现：

cond = (a > value1) & (b > value1)
A = a[cond]
B = b[cond]

UnkwnTech · Answer

您想要创建一个函数来执行上述操作，即如果数组中的元素大于某个特定值，则将其作为数组返回。

import numpy as np

def select_greater_elements(array1, array2, value1):
    condition = (array1 > value1) & (array2 > value1)
    A = array1[condition]
    B = array2[condition]
    return [A, B]

a = np.array([1, 2, 4, 3, 4, 7, 8])
b = np.array([6, 5, 2, 5, 7, 8, 9])

result = select_greater_elements(a, b, 2)
A = result[0]
B = result[1]
print(A)
print(B)

# 输出:
# [3, 4, 7, 8]
# [5, 7, 8, 9]

Wayne Conrad · Answer

通过使用布尔条件(array1 > value1) & (array2 > value1)创建了一个条件掩码，随后直接利用此条件掩码对array1和array2进行布尔索引，从而选取元素。这样做避免了显式追加的需要。

以下是改进后的代码示例：

import numpy as np

def 筛选数组(array1, array2, value1):
    条件 = (array1 > value1) & (array2 > value1)
    return array1[条件], array2[条件]

value1 = 2
a = np.array([1, 2, 4, 3, 4, 7, 8])
b = np.array([6, 5, 2, 5, 7, 8, 9])
A, B = 筛选数组(a, b, value1)
print(A)
print(B)
# 输出:
# [3 4 7 8]
# [5 7 8 9]

或者更简洁的方式：

a = np.array([1, 2, 4, 3, 4, 7, 8])
b = np.array([6, 5, 2, 5, 7, 8, 9])
value1 = 2  # 需要先定义value1变量

掩码 = (a > value1) & (b > value1)
A = a[掩码]
B = b[掩码]

print(A)
print(B)
# 同样输出:
# [3 4 7 8]
# [5 7 8 9]