如何在3D空间中为每个点查找最多6个分方向最近邻？

问题需求 #

需要在3D空间中为每个点查找最多6个最近邻，每个邻居对应左、右、前、后、上、下六个方向中的一个（角落点3个、边缘点4个、近壁点5个、内部点6个）。原有的双重循环分组方案处理70k点耗时约250分钟，效率极低；尝试过KDTree但点分布不均时会返回同一方向的多个邻居，不符合方向限制要求。

原示例数据 #

points = [
    (0, 0, 0), (0, 0, 1), (0, 0, 2),
    (0, 1, 0), (0, 1, 1), (0, 1, 2),
    (0, 2, 0), (0, 2, 1), (0, 2, 2),
    (1, 0, 0), (1, 0, 1), (1, 0, 2),
    (1, 1, 0), (1, 1, 1), (1, 1, 2),
    (1, 2, 0), (1, 2, 1), (1, 2, 2),
    (2, 0, 0), (2, 0, 1), (2, 0, 2),
    (2, 1, 0), (2, 1, 1), (2, 1, 2),
    (2, 2, 0), (2, 2, 1), (2, 2, 2)
]

原低效代码 #

for i, (x1, y1, z1) in enumerate(points):
    f_xy = lambda x, y: (x - x1) + (y - y1)
    g_xy = lambda x, y: (x - x1) - (y - y1)
    f_xz = lambda x, z: (x - x1) + (z - z1)
    g_xz = lambda x, z: (x - x1) - (z - z1)
    f_yz = lambda y, z: (y - y1) + (z - z1)
    g_yz = lambda y, z: (y - y1) - (z - z1)

    groups = {'right': [], 'left': [], 'front': [], 'back': [], 'top': [], 'bottom': []}

    for j, (x2, y2, z2) in enumerate(points):
        if i != j:
            if f_xy(x2, y2) >= 0 and g_xy(x2, y2) >= 0 and f_xz(x2, z2) >= 0 and g_xz(x2, z2) >= 0: groups['right'].append((x2, y2, z2))
            if f_xy(x2, y2) <= 0 and g_xy(x2, y2) <= 0 and f_xz(x2, z2) <= 0 and g_xz(x2, z2) <= 0: groups['left'].append((x2, y2, z2))

            if f_xy(x2, y2) > 0 and g_xy(x2, y2) < 0 and f_yz(y2, z2) >= 0 and g_yz(y2, z2) >= 0: groups['front'].append((x2, y2, z2))
            if f_xy(x2, y2) < 0 and g_xy(x2, y2) > 0 and f_yz(y2, z2) <= 0 and g_yz(y2, z2) <= 0: groups['back'].append((x2, y2, z2))

            if f_xz(x2, z2) < 0 and g_xz(x2, z2) > 0 and f_yz(y2, z2) < 0 and g_yz(y2, z2) > 0: groups['bottom'].append((x2, y2, z2))
            if f_xz(x2, z2) > 0 and g_xz(x2, z2) < 0 and f_yz(y2, z2) > 0 and g_yz(y2, z2) < 0: groups['top'].append((x2, y2, z2))

    print(f'Vertex: {(x1, y1, z1)}')
    print(len(groups['right']) + len(groups['left']) + len(groups['front']) + len(groups['back']) + len(groups['bottom']) + len(groups['top']))
    print('Left', groups['left'])
    print('Right', groups['right'])
    print('Front', groups['front'])
    print('Back', groups['back'])
    print('Top', groups['top'])
    print('Bottom', groups['bottom'])
    print('\n')

原方案效率瓶颈 #

原代码采用**O(n²)**的双重循环，对每个点遍历所有其他点进行方向判断，数据量增大时耗时呈指数增长，这是70k点处理缓慢的核心原因。

高效解决方案 #

核心思路是通过空间索引或哈希查询将时间复杂度降至线性或对数级别，同时严格过滤每个方向的最近点。

方案1：哈希集合精确匹配（适用于规则网格点） #

如果点是规则网格上的点（如整数坐标），直接计算每个方向的目标坐标，用集合快速查询是否存在：

points = [
    (0, 0, 0), (0, 0, 1), (0, 0, 2),
    (0, 1, 0), (0, 1, 1), (0, 1, 2),
    (0, 2, 0), (0, 2, 1), (0, 2, 2),
    (1, 0, 0), (1, 0, 1), (1, 0, 2),
    (1, 1, 0), (1, 1, 1), (1, 1, 2),
    (1, 2, 0), (1, 2, 1), (1, 2, 2),
    (2, 0, 0), (2, 0, 1), (2, 0, 2),
    (2, 1, 0), (2, 1, 1), (2, 1, 2),
    (2, 2, 0), (2, 2, 1), (2, 2, 2)
]

# 转换为集合，实现O(1)查询
point_set = set(points)

# 定义六个方向的偏移量（可根据网格步长调整）
directions = {
    'right': (1, 0, 0),
    'left': (-1, 0, 0),
    'front': (0, 1, 0),
    'back': (0, -1, 0),
    'top': (0, 0, 1),
    'bottom': (0, 0, -1)
}

for point in points:
    x, y, z = point
    neighbors = {}
    for dir_name, (dx, dy, dz) in directions.items():
        neighbor = (x + dx, y + dy, z + dz)
        if neighbor in point_set:
            neighbors[dir_name] = neighbor
    # 输出结果
    print(f'Vertex: {point}')
    print(f'邻居数量: {len(neighbors)}')
    for dir_name, neighbor in neighbors.items():
        print(f'{dir_name}: {neighbor}')
    print('\n')

方案2：KDTree结合方向过滤（适用于任意分布点） #

对于非规则分布的点，用cKDTree快速查找近邻，再根据向量主导分量过滤出每个方向的最近点：

import numpy as np
from scipy.spatial import cKDTree

points = np.array([
    (0, 0, 0), (0, 0, 1), (0, 0, 2),
    (0, 1, 0), (0, 1, 1), (0, 1, 2),
    (0, 2, 0), (0, 2, 1), (0, 2, 2),
    (1, 0, 0), (1, 0, 1), (1, 0, 2),
    (1, 1, 0), (1, 1, 1), (1, 1, 2),
    (1, 2, 0), (1, 2, 1), (1, 2, 2),
    (2, 0, 0), (2, 0, 1), (2, 0, 2),
    (2, 1, 0), (2, 1, 1), (2, 1, 2),
    (2, 2, 0), (2, 2, 1), (2, 2, 2)
])

# 构建KDTree
tree = cKDTree(points)

# 根据向量主导分量判断方向
def get_direction(vec):
    x, y, z = vec
    max_abs = max(abs(x), abs(y), abs(z))
    if abs(x) == max_abs:
        return 'right' if x > 0 else 'left'
    elif abs(y) == max_abs:
        return 'front' if y > 0 else 'back'
    else:
        return 'top' if z > 0 else 'bottom'

# 为每个点查找方向最近邻
for i, point in enumerate(points):
    # 查找前10个近邻（确保覆盖所有可能方向）
    distances, indices = tree.query(point, k=10)
    # 跳过自身
    valid_mask = indices != i
    valid_indices = indices[valid_mask]
    valid_distances = distances[valid_mask]
    
    direction_neighbors = {}
    for idx, dist in zip(valid_indices, valid_distances):
        neighbor = points[idx]
        vec = neighbor - point
        dir_name = get_direction(vec)
        # 仅保留该方向最近的点
        if dir_name not in direction_neighbors or dist < direction_neighbors[dir_name][0]:
            direction_neighbors[dir_name] = (dist, neighbor)
    
    # 整理输出结果
    result = {dir: tuple(neigh) for dir, (dist, neigh) in direction_neighbors.items()}
    print(f'Vertex: {tuple(point)}')
    print(f'邻居数量: {len(result)}')
    for dir_name, neighbor in result.items():
        print(f'{dir_name}: {neighbor}')
    print('\n')

方案优势 #

• 方案1时间复杂度为O(n)，仅适用于规则网格点，速度最快；
• 方案2时间复杂度为O(n log n)，适用于任意分布的点，处理70k点仅需数秒；
• 两种方案都能严格保证每个方向最多一个最近邻，完全符合需求。

内容的提问来源于stack exchange，提问作者jan-swiatek