Python 如何优化numpy循环，该循环对由另一个数组索引的数组中的值求和，其中值等于循环索引

我有一段代码，在应用程序运行期间被多次调用。 它采用一个表示值的数字数组（value\u数组）。 这些应汇总到区域中，这些区域在区域数组中定义。 zone_ID表示zone_数组中所有可能区域的列表

基本上是这样的：我有一张人口栅格地图，我想知道每个分区地图上有多少人

守则：

values = np.zeros(len(zone_ids))
for i in zone_ids:
    values[i] = round(np.nansum(value_array[zone_array == i]), 2)
return values

罪魁祸首似乎是for循环，但我没有找到消除它的方法，也没有得到相同的结果

我试着用bincount，但没有成功。 使用numba jit也没有效果

我想远离cython，因为这段代码将用于Qgis插件中，该插件不支持cython

测试代码：

import numpy as np


def fill_values(zone_array, value_array, zone_ids):
    values = np.zeros(len(zone_ids))
    for i in zone_ids:
        values[i] = round(np.nansum(value_array[zone_array == i]), 2)
    return values


def run():
    # 300 different zones
    zone_ids = range(300)
    # zone map with 300 zones
    zone_array = (np.random.rand(2000, 2000) * 300).astype(int)
    # value map from which we want the sum of values per zone (real map can have NaN values)
    value_array = (np.random.rand(2000, 2000) * 10.)
    value_array[5, 5] = np.NAN
    fill_values(zone_array, value_array, zone_ids)


if __name__ == '__main__':
    run()

每个回路1.92 s±17.5 ms（7次运行的平均值±标准偏差，每个回路1次）

按照Divakar的建议实施bincount：

---

每个循环203 ms±15.2 ms（7次运行的平均值±标准偏差，每个循环1次）

如果直接使用

bincount

，则总和中会有

NaNs

。因此，您只需将

NaNs

替换为

zero

，并使用

bincount

。作为矢量化解决方案，这应该快得多

因此，实施将是必要的-

val_nonan = np.where(np.isnan(value_array), 0, value_array)
out = np.round(np.bincount(zone_array.ravel(), val_nonan.ravel()),2)

---

罪魁祸首不是for循环。相反，问题在于内部的比较

zone\u array==i

。必须检查每个分区的所有2000x2000=4e6值是否等于

i

。如果我减少分区id的数量，我将获得速度提升，因此for循环仍然涉及性能问题。因为我知道我没有选择不做

zone\u array==i

的方法，所以我专注于循环。最好的方法是我可以使用

zone\u array==zone\u ids

跳过循环。您可以广播与

zone\u array[：，：，None]==zone\u ids

的比较，但这仍然将索引保留在for循环中，并且不会给性能带来太多改善。这解决了我的问题。谢谢。我猜我的bincount会尝试被nan值搞砸的地方。另外，

values=out[zone_id]

用于需要分区子集结果的情况。