如何高效实现Numpy坐标点数组的计数累加？

需求说明 #

给定Numpy坐标点数组，需生成计数数组：每个坐标点的出现次数，对应结果数组中该位置的数值。例如坐标[1,1,1]出现2次，对应位置值就为2。

示例输入 #

import numpy as np
points = np.array([[1,1,1], [0,1,1], [1,0,0], [1,1,1]])

示例输出结构 #

r = [[[0,0],
      [0,1]],
     [[1,0],
      [0,2]]]

原生Python循环实现（速度过慢） #

# 假设r已初始化为对应形状的零数组
for p in points:
    r[p[0]][p[1]][p[2]] += 1

该需求用于实时场景：每帧最多30万个点，结果矩阵最大为400x400x400，且点数和矩阵形状动态变化，需基于Numpy实现高效方案。

高效实现方案 #

方案1：利用np.unique批量赋值 #

通过np.unique提取唯一坐标及对应计数，再批量写入结果数组，避免Python循环开销：

def count_points(points, output_shape):
    # 初始化结果数组为0
    r = np.zeros(output_shape, dtype=np.int32)
    # 获取唯一坐标和对应出现次数
    unique_points, counts = np.unique(points, axis=0, return_counts=True)
    # 批量给对应坐标位置赋值计数
    r[tuple(unique_points.T)] = counts
    return r

# 调用示例
output_shape = (2, 2, 2)
r = count_points(points, output_shape)

此方案适合坐标重复率较高的场景，内存占用更低。

方案2：利用np.bincount（高吞吐量场景首选） #

将多维坐标转换为一维索引，用np.bincount统计次数后重塑为目标形状，底层实现高度优化，适合30万级别的实时处理：

def count_points_bincount(points, output_shape):
    # 计算一维索引的步长
    strides = np.array([output_shape[1]*output_shape[2], output_shape[2], 1])
    # 将三维坐标转为一维索引
    indices = points.dot(strides)
    # 统计每个索引的出现次数，确保长度匹配结果矩阵总元素数
    counts = np.bincount(indices, minlength=np.prod(output_shape))
    # 重塑为目标三维形状
    return counts.reshape(output_shape)

# 调用示例
output_shape = (2, 2, 2)
r = count_points_bincount(points, output_shape)

方案对比 #

• 坐标重复率高时，np.unique方案内存占用更低；
• 点数多、坐标分布均匀时，np.bincount方案速度更快，时间复杂度接近O(n)。

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