带逆索引的快速体素化方法优化问询

方案1：将三维体素索引编码为一维整数（最快的numpy原生方案） #

np.unique对一维数组的处理效率远高于二维数组，我们可以把三维体素索引(x,y,z)编码为唯一的一维整数，借助一维np.unique获取唯一值和逆索引，最后解码回三维格式。

代码实现 #

import numpy as np

def voxelize_fast(points, voxel_size):
    voxel_indices = np.floor(points / voxel_size).astype(np.int32)
    
    # 编码三维索引为一维：选择足够大的基数避免冲突（可根据实际场景调整）
    base = np.array([1, 2**20, 2**40], dtype=np.int64)
    voxel_1d = voxel_indices.dot(base)
    
    # 一维数组的np.unique性能远优于二维场景
    unique_1d, inverse_indices = np.unique(voxel_1d, return_inverse=True)
    
    # 解码回三维唯一体素索引
    unique_voxel_indices = np.column_stack([
        unique_1d // base[0] % base[1],
        unique_1d // base[1] % base[2],
        unique_1d // base[2]
    ]).astype(np.int32)
    
    voxel_centers = (unique_voxel_indices + 0.5) * voxel_size
    return voxel_indices, unique_voxel_indices, voxel_centers, inverse_indices

说明 #

• 编码基数需根据体素索引的最大范围调整，确保不同三维索引不会被编码为同一整数；若点云范围有限，可先将体素索引平移至非负区间再编码。
• 完全基于numpy原生操作，无额外依赖，性能比原二维np.unique方案提升数倍。

方案2：Numba手动实现二维Unique带逆索引 #

若需进一步加速，可借助Numba手动实现二维数组的唯一值提取与逆索引生成，绕过numpy在二维场景的性能开销。

代码实现 #

import numpy as np
from numba import njit

@njit
def numba_unique_2d(arr, return_inverse=True):
    # 排序数组并保留原始索引
    sorted_indices = np.argsort(arr.view(np.int32).ravel())
    sorted_arr = arr[sorted_indices]
    
    # 标记唯一元素的位置
    unique_mask = np.zeros(len(sorted_arr), dtype=np.bool_)
    unique_mask[0] = True
    for i in range(1, len(sorted_arr)):
        if not np.all(sorted_arr[i] == sorted_arr[i-1]):
            unique_mask[i] = True
    unique_voxels = sorted_arr[unique_mask]
    
    if not return_inverse:
        return unique_voxels
    
    # 生成逆索引：映射原始元素到唯一值索引
    inverse = np.zeros(len(arr), dtype=np.int32)
    current_idx = 0
    inverse[sorted_indices[0]] = current_idx
    for i in range(1, len(sorted_arr)):
        if unique_mask[i]:
            current_idx += 1
        inverse[sorted_indices[i]] = current_idx
    
    return unique_voxels, inverse

def voxelize_numba(points, voxel_size):
    voxel_indices = np.floor(points / voxel_size).astype(np.int32)
    unique_voxel_indices, inverse_indices = numba_unique_2d(voxel_indices)
    voxel_centers = (unique_voxel_indices + 0.5) * voxel_size
    return voxel_indices, unique_voxel_indices, voxel_centers, inverse_indices

说明 #

• 首次调用Numba函数存在编译开销，但后续数千次调用的性能会远高于numpy方案。
• 通过排序+遍历的方式定位唯一元素，逻辑清晰且高效。

方案3：哈希表映射（适合极端大/稀疏场景） #

若体素索引范围极大，一维编码可能溢出，可借助Numba typed字典实现三维索引到唯一ID的映射，时间复杂度接近O(n)。

代码实现 #

import numpy as np
from numba import njit, types
from numba.typed import Dict

@njit
def hash_voxel_indices(voxel_indices):
    voxel_dict = Dict.empty(
        key_type=types.UniTuple(types.int32, 3),
        value_type=types.int32
    )
    inverse_indices = np.zeros(len(voxel_indices), dtype=np.int32)
    current_id = 0
    
    for i in range(len(voxel_indices)):
        key = (voxel_indices[i,0], voxel_indices[i,1], voxel_indices[i,2])
        if key not in voxel_dict:
            voxel_dict[key] = current_id
            current_id += 1
        inverse_indices[i] = voxel_dict[key]
    
    # 提取唯一体素索引
    unique_voxels = np.zeros((current_id, 3), dtype=np.int32)
    for key, idx in voxel_dict.items():
        unique_voxels[idx] = key
    
    return unique_voxels, inverse_indices

def voxelize_hash(points, voxel_size):
    voxel_indices = np.floor(points / voxel_size).astype(np.int32)
    unique_voxel_indices, inverse_indices = hash_voxel_indices(voxel_indices)
    voxel_centers = (unique_voxel_indices + 0.5) * voxel_size
    return voxel_indices, unique_voxel_indices, voxel_centers, inverse_indices

说明 #

• 哈希表方案在体素数量极多且分布稀疏时，性能优于排序类方案。
• 需通过Numba typed字典指定类型，确保高效运行。

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