如何无重复划分点云？PyTorch实现及分组方案咨询

关于reshape/view的直接使用问题 #

绝对不能用reshape或view实现需求。这两个方法仅仅是对张量的内存布局进行重新映射，完全不考虑点云的空间位置关系，只会把点按原始顺序硬拆成每组K个，和“最近邻分组”的要求完全不沾边。

普通K近邻重叠问题的原因 #

你之前用K近邻出现重叠，是因为常规K近邻是为每个点单独寻找邻居，每个点会被多个其他点当作邻居，自然会出现重复归属的情况。而我们需要的是划分式分组——每个点只能属于一个组，这本质是一种聚类任务，而非普通的邻域搜索。

可行的实现方案（PyTorch） #

核心思路是贪心式最近邻划分：先选中心点，再为每个中心点选取未被选中的最近邻点组成无重叠的组，具体步骤和代码如下：

步骤说明 #

• 维护一个标记数组，记录每个点是否已被分配到组中
• 循环采样未被标记的点作为中心点，为其选取K-1个最近的未标记点，组成一组
• 标记这些点为已选中，直到所有点都被分配

示例代码 #

import torch

def greedy_nearest_neighbor_grouping(points, K):
    B, N, C = points.shape
    assert N % K == 0, "点的数量N必须是K的整数倍"
    num_groups = N // K
    
    # 初始化已选中标记，(B, N)
    selected = torch.zeros(B, N, dtype=torch.bool, device=points.device)
    # 存储最终分组结果
    groups = torch.zeros(B, num_groups, K, C, device=points.device)
    
    for batch_idx in range(B):
        batch_points = points[batch_idx]
        batch_selected = selected[batch_idx]
        batch_groups = groups[batch_idx]
        
        for group_idx in range(num_groups):
            # 找到未被选中的点的索引
            unselected_indices = torch.where(~batch_selected)[0]
            # 随机选一个中心点（也可以用最远点采样FPS选中心点，提升分布均匀性）
            center_idx = unselected_indices[torch.randint(0, len(unselected_indices), (1,))]
            # 计算中心点到所有未选中点的距离
            center_point = batch_points[center_idx].unsqueeze(0)
            distances = torch.cdist(center_point, batch_points[unselected_indices])[0]
            # 取距离最近的K个点（包括中心点自己）
            topk_indices = torch.topk(distances, k=K, largest=False)[1]
            selected_indices = unselected_indices[topk_indices]
            # 将这些点存入分组
            batch_groups[group_idx] = batch_points[selected_indices]
            # 标记为已选中
            batch_selected[selected_indices] = True
    
    return groups

优化方向 #

如果觉得随机选中心点导致分组分布不够均匀，可以替换成**最远点采样(FPS)**来选择中心点：先选一个初始点，之后每次选离已选中心点最远的点作为下一个中心点，这样每个组的空间分布会更均衡，避免局部扎堆。

补充说明 #

• 必须保证N是K的整数倍，如果实际数据不满足，可以提前随机丢弃部分点，或者复制少量点补全（根据业务需求选择）
• 上述代码在GPU上可以高效运行，适合批量处理点云数据

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