# 如何使用if条件在1D和2D numpy数组之间进行向量化计算

``````import numpy as np

r_base = np.linspace(0, 4, 5)
np.random.seed(0)
r_mat = np.array([r_base * np.random.uniform(0.9, 1.1, 5),
r_base * np.random.uniform(0.9, 1.1, 5),
r_base * np.random.uniform(0.9, 1.1, 5)])

a_array = np.linspace(1, 3, 3)

def func_vetorized_level1(r_row, a):
if r_row.mean() > 2:
result = np.where((r_row >= a), r_row - a, np.nan)
else:
result = np.where((r_row >= a), r_row + a, 0)
return result
# try to broadcast this func to every row of r_mat using list comprehension
res_mat = np.array([func_vetorized_level1(this_r_row, this_a)
for this_r_row, this_a in zip(r_mat, a_array)])
``````

``````res_mat =
array([[       nan, 0.04303787, 1.04110535, 2.02692991, 2.93892384],
[       nan,        nan, 0.1567092 , 1.27819766, 1.90675322],
[0.        , 0.        , 0.        , 6.25535798, 6.65682885]])
``````

``````r_base = np.arange(5.)
a_array = np.arange(1., 4.)
``````

``````np.random.seed(0)
r_mat = r_base * np.random.uniform(0.9, 1.1, (3, 5))
``````

``````out = np.empty_like(r_mat)
``````

``````a = a_array[:, None]
``````

``````row_mask = (r_mat.mean(axis=1) > 2)[:, None]
``````

``````np.subtract(r_mat, a, out=out, where=row_mask & elem_mask)
``````

0条评论