如何用Numpy优化需求函数以消除嵌套循环？

原函数的核心逻辑是：对输入的每个价格值，统计输入价格列表中大于等于该价格的元素数量，得到对应需求人数数组。用Numpy的向量化操作可以彻底去掉嵌套循环，大幅提升效率，尤其是当价格数据量较大时。

优化后的代码 #

import numpy as np

def demand_y(self, prices):
    prices_np = np.array(prices)
    # 向量化比较+行求和，直接得到每个价格对应的符合条件的人数
    y = (prices_np >= prices_np[:, np.newaxis]).sum(axis=1)
    return y.tolist()  # 若需保持原函数返回列表的格式，可转成Python列表

逻辑说明 #

1. 先将输入的prices转为Numpy数组，利用其底层C实现的向量化运算能力，规避Python循环的性能开销。
2. prices_np[:, np.newaxis]将一维数组转为列向量，和原数组做>=比较时，Numpy会自动广播生成n×n的布尔数组（n为价格数量），每个元素(i,j)代表prices_np[i] >= prices_np[j]。
3. 对布尔数组按行求和（sum(axis=1)），布尔值会自动转为0/1，求和结果就是对应价格下愿意购买的人数。
4. 若需要和原函数返回格式一致（Python列表），用.tolist()转换即可。

效率优势 #

原函数是O(n²)的嵌套Python循环，Numpy的向量化操作无Python层面的循环开销，当价格数据量较大（如上万条）时，运行速度会有几十甚至上百倍的提升。

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