如何用Python实现numpy整数数组的按条件乘除运算

直接用Numpy原生向量化操作实现，无Python层循环，处理大数组性能优异。

完整代码 #

import numpy as np

def process_arr(input_arr: np.ndarray) -> np.ndarray:
    # 校验输入是否为整数类型数组
    if not np.issubdtype(input_arr.dtype, np.integer):
        raise TypeError("输入必须为整数类型的numpy数组")
    # 生成可被7整除的元素掩码
    mask = input_arr % 7 == 0
    # 初始化结果：默认所有元素乘3
    res = input_arr * 3
    # 替换符合条件的元素为除以7的结果
    res[mask] = input_arr[mask] // 7
    return res

测试示例 #

# 测试输入
test_arr = np.array([7, 2, 14, 9, 21, 5, 0], dtype=np.int64)
print(process_arr(test_arr))
# 输出：[ 1  6  2 27  3 15  0]

说明 #

• 提前做了输入类型校验，避免非整数数组传入后得到不符合预期的结果
• 全程使用Numpy向量化操作，时间复杂度为O(n)，处理十万级以上元素的数组时，效率比Python原生循环高几十到上百倍
• 默认使用整数除法//保证输出仍为整数类型，若需要返回浮点结果，把//替换为/即可
• 0可以被7整除，会被处理为0//7=0，符合数学规则

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