Python中高效计算网格单元内观测值均值的最优方法探究

我有一个二维数据点数组（X）以及对应的观测值（z），想要计算每个网格单元内z值的均值网格。用NumPy嵌套for循环的方式效率极低，有没有用内置函数或列表推导式的更快实现方法？尽量不想用Numba/jit。

原始低效实现（一维坐标版本） #

import numpy

x = numpy.random.rand(1000000)
y = numpy.random.rand(1000000)
z = numpy.random.rand(1000000)

nx = 1000
ny = 1000
xl = numpy.linspace(0,1,nx+1)
yl = numpy.linspace(0,1,ny+1)
zm = numpy.full((nx,ny),numpy.nan)
for i in range(nx):
    for j in range(ny):
        zm[i,j] = numpy.mean(z, where = ((x>xl[i]) & (x<=xl[i+1]) & (y>yl[j]) & (y<=yl[j+1]))) # 单次循环约4.5ms，总耗时约75分钟

原始低效实现（二维数组版本） #

import numpy

X = numpy.array([numpy.random.rand(1000000),numpy.random.rand(1000000)]).T
z = numpy.random.rand(1000000)

nx = 1000
ny = 1000
xl = numpy.linspace(0,1,nx+1)
yl = numpy.linspace(0,1,ny+1)
zm = numpy.full((nx,ny),numpy.nan)
for i in range(nx):
    print(i)
    for j in range(ny):
        zm[i,j] = numpy.mean(z, where = ((X[:,0]>xl[i]) & (X[:,0]<=xl[i+1]) & (X[:,1]>yl[j]) & (X[:,1]<=yl[j+1]))) # 单次循环约4.5ms，总耗时约75分钟

高效实现方案（完全向量化，无循环） #

方法1：使用numpy.histogram2d #

利用histogram2d同时计算每个网格的点数和z值总和，再通过除法得到均值。该方法完全向量化，时间复杂度为O(N)（N为数据点数量），相比嵌套循环的O(nxnyN)效率提升几个数量级。

import numpy

x = numpy.random.rand(1000000)
y = numpy.random.rand(1000000)
z = numpy.random.rand(1000000)

nx = 1000
ny = 1000
xl = numpy.linspace(0, 1, nx+1)
yl = numpy.linspace(0, 1, ny+1)

# 计算每个网格的点数和z值总和
counts, _, _ = numpy.histogram2d(x, y, bins=[xl, yl])
z_sums, _, _ = numpy.histogram2d(x, y, bins=[xl, yl], weights=z)

# 计算均值，空网格自动得到nan（0/0在numpy中返回nan）
zm = z_sums / counts

方法2：使用numpy.digitize+numpy.bincount #

先通过digitize确定每个点所属的网格索引，再用bincount统计每个网格的点数和z总和，最后计算均值。同样是全向量化操作，适合需要更细粒度控制索引的场景。

import numpy

x = numpy.random.rand(1000000)
y = numpy.random.rand(1000000)
z = numpy.random.rand(1000000)

nx = 1000
ny = 1000
xl = numpy.linspace(0, 1, nx+1)
yl = numpy.linspace(0, 1, ny+1)

# 获取每个点对应的网格索引（0-based）
i = numpy.digitize(x, xl) - 1
j = numpy.digitize(y, yl) - 1

# 将二维索引转换为一维，方便bincount统计
flat_idx = i * ny + j

# 统计每个网格的点数和z总和
counts = numpy.bincount(flat_idx, minlength=nx*ny).reshape(nx, ny)
z_sums = numpy.bincount(flat_idx, weights=z, minlength=nx*ny).reshape(nx, ny)

# 计算均值，空网格为nan
zm = z_sums / counts

注意事项 #

• 两种方法都自动处理空网格（无数据点的网格），返回nan，和原始代码逻辑一致。
• 对于二维数组X的情况，只需将x替换为X[:,0]，y替换为X[:,1]即可，代码逻辑完全通用。
• 实际测试中，这两种方法处理1e6个数据点、1000x1000网格的耗时仅需几十毫秒，远优于原始循环的75分钟。

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