如何加速Python中基于cKDTree的近邻连线计算？

Alright, let's tackle this performance boost head-on. Your original approach (I’m guessing brute-force checking every point pair) would crawl to a halt with 5k-10k points—scipy.spatial.cKDTree is the perfect tool here because it turns an O(n²) problem into something manageable, with O(n log n) preprocessing and O(log n) per query.

Here's the Optimized Code #

First, make sure you have scipy installed (pip install scipy if you don't). Then drop this in place of your current function:

import numpy as np
from scipy.spatial import cKDTree

def calculate_draw_vertice_lines(vertices_xy_list, max_distance=55, max_neighbors=3):
    # Convert your nested list to a numpy array—cKDTree works far more efficiently with these
    points = np.array(vertices_xy_list)
    
    # Build the spatial index once (this is the O(n log n) heavy-lifting step)
    point_tree = cKDTree(points)
    
    lines = []
    
    for idx, current_point in enumerate(points):
        # Query for the closest max_neighbors +1 points (the +1 skips the point itself)
        # We set a distance upper bound to ignore anything beyond 55 pixels immediately
        distances, neighbor_indices = point_tree.query(
            current_point, 
            k=max_neighbors + 1, 
            distance_upper_bound=max_distance
        )
        
        # Filter out self-references and points that are outside the 55-pixel limit
        valid_neighbors = [
            neighbor_idx 
            for neighbor_idx, dist in zip(neighbor_indices, distances)
            if neighbor_idx != idx and dist != np.inf
        ]
        
        # Add lines, but avoid duplicates (A->B and B->A are the same line for drawing)
        for neighbor_idx in valid_neighbors:
            if idx < neighbor_idx:
                lines.append([current_point.tolist(), points[neighbor_idx].tolist()])
    
    return lines

What Makes This Faster? #

• Spatial Indexing: Building the cKDTree organizes your points into a hierarchical structure, so the algorithm can skip entire chunks of points when searching for neighbors—no need to check every single point against every other.
• Targeted Queries: The query() method directly grabs the closest N points within your distance limit, eliminating extra loops and redundant calculations.
• Duplicate Prevention: The idx < neighbor_idx check cuts your line count in half by only storing each unique line once, avoiding redundant entries that would slow down downstream drawing logic.

Quick Performance Note #

If you ever scale beyond 10k points, you could also use point_tree.query_ball_point() with k=max_neighbors to fetch closest points in the radius, but query() is simpler when you need a fixed maximum number of neighbors. Either way, both approaches are orders of magnitude faster than brute-force.

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