超大规模3D点云数组np.unique性能优化：高效替代方案咨询

针对你处理3600万条3D点数据时np.unique耗时过长的问题，以下是几个能显著提升速度的替代方案：

方案1：基于空间哈希+np.bincount的栅格计数 #

你的核心需求是按栅格（步长3）划分点云并统计每个栅格的点数，完全可以跳过全局排序的np.unique，直接用空间哈希生成唯一栅格键，再用np.bincount做线性时间统计：

import numpy as np
import time

points = np.random.random(size=(38867362, 3)) * 10000

start_time = time.time()
# 计算每个点的栅格坐标（整数化）
min_vals = np.min(points, axis=0)
grid_coords = ((points - min_vals) // 3).astype(np.int64)

# 将三维栅格坐标转换为一维唯一键（避免溢出，用int64）
max_grid = np.max(grid_coords, axis=0) + 1
grid_keys = np.ravel_multi_index(grid_coords.T, max_grid)

# 统计每个键的出现次数
counts = np.bincount(grid_keys)
# 获取非零计数的栅格键和对应坐标
non_zero_idx = counts > 0
unique_grid_keys = np.nonzero(non_zero_idx)[0]
unique_points = np.unravel_index(unique_grid_keys, max_grid)
unique_points = np.array(unique_points).T * 3 + min_vals  # 还原为栅格原点坐标
# 生成inverse数组（每个点对应的唯一索引）
inverse = np.searchsorted(unique_grid_keys, grid_keys)

print("::INFO:: Total time taken: ", time.time() - start_time)

为什么更快：np.bincount是O(n)时间复杂度，而np.unique(axis=0)需要先做全局排序（O(n log n)），对于超大规模数据，时间差距会非常明显。

方案2：用numba加速自定义分组逻辑 #

如果需要保留类似np.unique的输出格式，用numba对分组统计逻辑做JIT编译，利用CPU并行优化：

import numpy as np
import numba
import time

points = np.random.random(size=(38867362, 3)) * 10000

@numba.njit(parallel=True)
def count_unique_grids(grid_coords):
    # 先排序栅格坐标
    sorted_idx = np.argsort(grid_coords.view(np.int64, 3))
    sorted_coords = grid_coords[sorted_idx]
    unique_coords = []
    counts = []
    current = sorted_coords[0]
    count = 1
    for i in numba.prange(1, len(sorted_coords)):
        if np.all(sorted_coords[i] == current):
            count +=1
        else:
            unique_coords.append(current)
            counts.append(count)
            current = sorted_coords[i]
            count =1
    unique_coords.append(current)
    counts.append(count)
    # 生成inverse数组
    inverse = np.zeros(len(grid_coords), dtype=np.int64)
    current_idx =0
    for i in range(len(sorted_idx)):
        inverse[sorted_idx[i]] = current_idx
        if i < len(sorted_coords)-1 and not np.all(sorted_coords[i] == sorted_coords[i+1]):
            current_idx +=1
    return np.array(unique_coords), np.array(counts), inverse

start_time = time.time()
min_vals = np.min(points, axis=0)
grid_coords = ((points - min_vals) //3).astype(np.int64)
unique_points, counts, inverse = count_unique_grids(grid_coords)
# 还原为原始坐标范围
unique_points = unique_points *3 + min_vals
print("::INFO:: Total time taken: ", time.time() - start_time)

为什么更快：numba的JIT编译消除了numpy的内部开销，并行处理进一步利用多核CPU，排序后的线性遍历比np.unique的通用实现更高效。

方案3：改用pandas的分组统计 #

pandas的groupby内部有针对大数组的优化，适合快速分组计数：

import numpy as np
import pandas as pd
import time

points = np.random.random(size=(38867362, 3)) * 10000

start_time = time.time()
min_vals = np.min(points, axis=0)
grid_coords = ((points - min_vals) //3).astype(np.int64)
# 转换为DataFrame
df = pd.DataFrame(grid_coords, columns=['x','y','z'])
# 分组计数
grouped = df.groupby(['x','y','z']).size()
# 提取结果
unique_points = grouped.index.to_numpy() *3 + min_vals
counts = grouped.values
# 生成inverse数组
df['idx'] = df.groupby(['x','y','z']).ngroup()
inverse = df['idx'].values

print("::INFO:: Total time taken: ", time.time() - start_time)

注意：pandas会占用更多内存，若内存紧张优先考虑方案1。

额外优化建议 #

• 确保使用64位整数存储栅格键，避免溢出问题；
• 如果不需要inverse数组，方案1的速度还能再提升（跳过np.searchsorted）；
• 若系统有GPU，可尝试用CuPy替代numpy，GPU加速的排序和统计能带来数量级的提升。

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