如何用np.select调用不同函数修改全局变量？求优化方案

核心问题分析 #

你当前代码的关键bug：choices里直接写PMultipleDMultiple()会无条件执行所有四个函数，不管conditions是否匹配——这就是全局变量数据混乱的根源：所有函数的global变量修改逻辑都跑了一遍，互相覆盖了数据。

问题1：如何用np.select定义多个变量 #

np.select本身返回单个数组（或标量），要同时定义多个变量，有两种可行方案：

方案1：每个变量单独用np.select #

针对每个需要定义的变量，单独编写conditions和对应的choices数组，适合变量逻辑独立、可直接赋值的场景：

# 示例：定义var1、var2两个变量
var1_choices = [val1_case1, val1_case2, val1_case3, val1_case4]
var2_choices = [val2_case1, val2_case2, val2_case3, val2_case4]

var1 = np.select(conditions, var1_choices, default=np.nan)
var2 = np.select(conditions, var2_choices, default=np.nan)

方案2：让选择项返回多变量元组，再拆分 #

如果每个条件下的变量需要复杂计算，可以让每个处理函数返回元组，通过np.select匹配后拆分赋值：

def case1():
    return calc_var1_case1(), calc_var2_case1(), calc_var3_case1()
def case2():
    return calc_var1_case2(), calc_var2_case2(), calc_var3_case2()

# 注意：这里传入函数对象，不要加括号调用
choices = [case1, case2, case3, case4]
selected_func = np.select(conditions, choices, default=None)

# 单条数据（标量条件）场景
if selected_func:
    var1, var2, var3 = selected_func()
# 批量数据（数组条件）场景，需循环处理每个元素
for func in selected_func:
    if func:
        v1, v2, v3 = func()
        # 批量存储逻辑

问题2：比np.select更快的方法 #

速度取决于你的数据类型（标量/数组）：

• 标量条件场景：直接用if-elif链比np.select更快——numpy函数有额外的数组处理开销，标量判断用原生分支更高效：

  if (PCount > 1) & (DCount > 1):
      var1, var2 = case1()
  elif (PCount > 1) & (DCount == 1):
      var1, var2 = case2()
  # 后续条件分支...
  

• 数组条件场景：用矢量化赋值替代np.select，可读性更好且速度相当：

  var1 = np.full_like(PCount, np.nan)
  var1[(PCount >1) & (DCount>1)] = val1_case1
  var1[(PCount >1) & (DCount==1)] = val1_case2
  

问题3：基于条件触发函数并安全处理全局变量 #

正确触发函数的方式 #

不要在choices里直接调用函数，而是传入函数对象，先匹配条件选出函数，再执行：

conditions = [
    ((PCount > 1) & (DCount > 1)),
    ((PCount > 1) & (DCount == 1)),
    ((PCount == 1) & (DCount > 1)),
    ((PCount == 1) & (DCount == 1)),
]

# 传入函数对象，不带括号
choices = [
    PMultipleDMultiple,
    PMultipleDSingle,
    PSingleDMultiple,
    PSingleDSingle,
]

# 获取匹配的函数
selected_func = np.select(conditions, choices, default=None)

# 仅执行匹配的函数
if selected_func is not None:
    selected_func()

避免全局变量的最佳实践 #

全局变量容易引发数据混乱，建议让函数返回需要的变量，在主逻辑中接收：

# 改造函数，去掉global声明，直接返回变量
def PMultipleDMultiple():
    var_a = 100
    var_b = "foo"
    var_c = np.array([1,2,3])
    return var_a, var_b, var_c

# 主逻辑中接收变量
if (PCount >1) & (DCount>1):
    var_a, var_b, var_c = PMultipleDMultiple()

如果必须使用全局变量，务必通过上述“先选函数再执行”的方式，确保只有匹配条件的函数才会修改全局变量，避免所有函数都执行导致的覆盖问题。

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