如何基于值提取defaultdict的键并绘图及优化大循环效率

原代码的核心性能瓶颈 #

• 双重for循环遍历全范围，哪怕ddtype中不存在的键也会触发查询生成默认值，产生大量无效计算
• 每次符合条件就单独调用plt.plot()，matplotlib的单次绘图调用开销很高，循环中反复调用会极大拖慢运行速度

最优优化方案（仅遍历已有键） #

如果你的需求只需要处理ddtype中已经存在的键值对，直接遍历字典条目批量收集坐标，仅调用一次绘图接口即可，性能提升最为明显：

import matplotlib.pyplot as plt
from collections import defaultdict

ddtype = defaultdict(int, {(1,1):2.2,(1,2):0.7,(2,1):0.9,(2,2):1.3})

x = []
y = []
for (i,j), val in ddtype.items():
    if val > 1.5:
        x.append(i)
        y.append(j)

# 单次绘图完成所有点的渲染
plt.scatter(x, y, marker='+', c='k')
plt.show()

需全范围遍历场景的优化方案 #

如果业务逻辑要求必须检查指定范围内的所有(i,j)组合（包括ddtype中不存在的键），可以用numpy向量化操作替代Python层循环，减少循环开销：

import matplotlib.pyplot as plt
import numpy as np
from collections import defaultdict

ddtype = defaultdict(int, {(1,1):2.2,(1,2):0.7,(2,1):0.9,(2,2):1.3})

# 定义需要遍历的i、j范围
i_start, i_end = 1, 3
j_start, j_end = 1, 3

# 生成坐标网格
i_arr, j_arr = np.meshgrid(np.arange(i_start, i_end), np.arange(j_start, j_end), indexing='ij')
# 批量判断符合条件的点
mask = np.vectorize(lambda i,j: ddtype[i,j] > 1.5)(i_arr, j_arr)

# 单次绘图
plt.scatter(i_arr[mask], j_arr[mask], marker='+', c='k')
plt.show()

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