Python函数处理单个数值正常，处理Numpy数组时结果异常

作为Python编程新手，我编写了一个基于a、b、c相对位置计算的函数，单个数值测试时结果正常：

a = 6
b = 3
c = 2

def function(a, b, c):
    if ((a >= b) & (b >= c)):
        return b - c 
    elif ((a <= b) & (b >= c)):
        return a - c
    else:
        return 0 
t = function(a, b, c)
print(t)  # 返回1，结果正确

但当a、b、c为Numpy数组时，整个结果数组只会返回b - c的结果，数组版本的函数代码如下：

def function(a, b, c):
    if ((a >= b) & (b >= c)).any():
        return b - c 
    elif ((a <= b) & (b >= c)).any():
        return a - c
    else:
        return 0 
t = function(a, b, c[i>1])
print(t)

（注：[i>1]是因数组输入数量可变，i=0时将使用另一个函数）

我还尝试使用np.where，但结果仍相同：

t = np.where(((prev2 >= Head_ELV) & (Head_ELV >= Bottom_ELV)).any, Head_ELV - Bottom_ELV, 0)

请问使用while循环是否能解决该问题？

你代码的核心问题出在.any()的使用上：.any()会检查整个数组中是否存在任意一个元素满足条件，只要有一个元素符合，就会直接返回整个b - c或a - c数组，完全不会逐元素判断每个位置的条件。比如第一个if里只要数组中有一个位置满足a>=b且b>=c，函数就直接返回b-c，忽略其他位置的情况。

另外你尝试的np.where代码里，.any是取了方法对象而非调用方法（应该写.any()），但就算调用了，依然是全局判断，无法实现逐元素的条件分支。

完全没必要用while循环——numpy的设计初衷就是用矢量化操作替代循环，循环不仅代码繁琐，效率也极低。推荐两种矢量化的解决方式：

方法一：嵌套使用np.where #

利用np.where的嵌套实现多条件分支，逐元素判断并赋值：

import numpy as np

def function(a, b, c):
    # 定义两个条件的掩码数组
    mask1 = (a >= b) & (b >= c)
    mask2 = (a <= b) & (b >= c)
    # 先处理mask1的情况，再处理mask2，剩余位置返回0
    result = np.where(mask1, b - c, 
                      np.where(mask2, a - c, 0))
    return result

方法二：布尔索引赋值 #

先初始化全0数组，再通过布尔索引给满足条件的位置赋值，逻辑更直观：

import numpy as np

def function(a, b, c):
    # 创建和输入数组同形状的全0结果数组
    result = np.zeros_like(a)
    mask1 = (a >= b) & (b >= c)
    mask2 = (a <= b) & (b >= c)
    # 给满足条件的位置分别赋值
    result[mask1] = b[mask1] - c[mask1]
    result[mask2] = a[mask2] - c[mask2]
    return result

这两种方法都会对数组中的每个元素单独判断条件，返回符合预期的结果数组。

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