如何修改Numpy三元运算逻辑实现按数组来源乘对应系数

原代码错误点 #

• 直接用==对比numpy数组会返回布尔数组，无法作为三元判断的布尔条件，需用np.all()判断数组是否全部元素相等
• 用逗号分隔两个数组切片会得到元组，而非合并后的单个数组，需要做扁平化拼接处理
• 三元运算优先级错误，原写法判断条件仅作用于第二个切片表达式，逻辑不符合预期

修正后的代码 #

写法1：批量计算（适合大数组，性能更高） #

import numpy as np 

a_multiplier = 3
b_multiplier = 5

a = np.array([5,32,1,4])
b = np.array([1,5,11,3])
a_and_b = np.array([5,1,32,5,1,11,4,3])

a_and_b = (np.column_stack((a_and_b[::2] * a_multiplier, a_and_b[1::2] * b_multiplier)).ravel()
           if np.all(a_and_b[::2] == a) 
           else np.column_stack((a_and_b[::2] * b_multiplier, a_and_b[1::2] * a_multiplier)).ravel())

print(a_and_b)

运行输出：

[15  5 96 25  3 55 12 15]

写法2：逐元素按索引判断（逻辑更直观，贴合需求） #

如果你需要直接通过元素索引判断来源，可以用逐元素三元判断的写法：

import numpy as np 

a_multiplier = 3
b_multiplier = 5

a_and_b = np.array([5,1,32,5,1,11,4,3])
a_and_b = np.array([val * a_multiplier if idx % 2 == 0 else val * b_multiplier for idx, val in enumerate(a_and_b)])

print(a_and_b)

运行结果和上述写法一致。

内容的提问来源于stack exchange，提问作者georgehere