如何用字典替代多分支if语句实现等高线图数据条件赋值？

问题背景 #

需要生成带多限制条件的等高线数据，原代码通过大量if-elif判断实现径向递减的z值，但希望用字典（键为条件阈值，值为对应数据系数）简化逻辑，避免冗余的条件判断。

原冗余代码：

import numpy as np

x = np.linspace(-50 , 50, 100)
y = np.linspace(-50 , 50, 100)
z = []
z_0 = 100
for i in x:
   for j in y:
      if i**2 + j**2 <= 10**2:
         z.append(1.9 * z_0)
      elif i**2 + j**2 <= 20**2:
         z.append(1.5 * z_0)
      elif i**2 + j**2 <= 30**2:
         z.append(1.4 * z_0)
      elif i**2 + j**2 <= 40**2:
         z.append(1.05 * z_0)
      else:
         z.append(z_0)

解决方案 #

方法1：基于字典的循环优化

核心是先对字典的阈值键从小到大排序，保证按原逻辑的优先级判断（最小阈值优先匹配），然后在循环中遍历排序后的键值对，找到第一个符合条件的系数：

import numpy as np

x = np.linspace(-50 , 50, 100)
y = np.linspace(-50 , 50, 100)
z = []
z_0 = 100
# 定义阈值-系数字典
requirements_dict = {10: 1.9, 20: 1.5, 30: 1.4, 40: 1.05}
# 对字典的键按从小到大排序，生成(阈值, 系数)的有序列表
sorted_thresholds = sorted(requirements_dict.items(), key=lambda item: item[0])

for i in x:
    for j in y:
        r_sq = i**2 + j**2
        # 遍历有序阈值，匹配第一个符合条件的系数
        for threshold, coeff in sorted_thresholds:
            if r_sq <= threshold**2:
                z.append(coeff * z_0)
                break
        else:
            # 所有阈值都不满足时，用默认值
            z.append(z_0)

方法2：Numpy向量化优化（更高效）

利用Numpy的向量化操作替代嵌套循环，大幅提升效率，同时结合字典的阈值逻辑：

import numpy as np

x = np.linspace(-50 , 50, 100)
y = np.linspace(-50 , 50, 100)
z_0 = 100
requirements_dict = {10: 1.9, 20: 1.5, 30: 1.4, 40: 1.05}

# 生成x-y网格矩阵
X, Y = np.meshgrid(x, y)
# 计算所有点的径向平方值
r_sq = X**2 + Y**2

# 对阈值从小到大排序，提取条件和对应系数
sorted_items = sorted(requirements_dict.items(), key=lambda item: item[0])
conditions = [r_sq <= threshold**2 for threshold, _ in sorted_items]
coeffs = [coeff for _, coeff in sorted_items]

# 使用np.select匹配条件，默认值为z0
z = np.select(conditions, [c * z_0 for c in coeffs], default=z_0)

说明 #

• 字典排序是关键：原代码的elif逻辑是从小到大匹配阈值，必须保证遍历顺序从最小阈值开始，否则会出现大阈值覆盖小阈值的错误。
• 向量化方法避免了Python级别的循环，处理大规模数据时效率远高于嵌套循环。

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