numpy中无循环获取聚类标注数组内各元素坐标的实现方法

纯NumPy向量化实现（无额外依赖） #

核心思路是先批量提取所有非零标签的坐标及对应标签值，再通过排序和分界点拆分得到每个标签的坐标组，全程无逐标签遍历数组的高开销循环，性能远高于逐标签调用np.where：

import numpy as np
from scipy.ndimage import label

# 原有标注逻辑
input_array = np.array([[0, 1, 1, 0],
                        [1, 1, 0, 0],
                        [0, 0, 0, 1],
                        [0, 0, 0, 1]])
labeled_array, num_labels = label(input_array)

# 无循环获取各标签坐标逻辑
# 1. 提取所有非零元素的坐标，形状为(N,2)，每行对应(行索引, 列索引)
coords = np.argwhere(labeled_array > 0)
# 2. 获取每个坐标对应的标签值
labels = labeled_array[coords[:, 0], coords[:, 1]]
# 3. 按标签值排序
sorted_idx = np.argsort(labels)
sorted_labels = labels[sorted_idx]
sorted_coords = coords[sorted_idx]
# 4. 定位不同标签的分界位置
split_points = np.where(np.diff(sorted_labels) != 0)[0] + 1
# 5. 拆分坐标组并转为目标字典格式
coord_groups = np.split(sorted_coords, split_points)
result = {label: group.tolist() for label, group in zip(np.unique(labels), coord_groups)}

运行后result输出和你要求的格式完全一致：

{1: [[0, 1], [0, 2], [1, 0], [1, 1]], 2: [[2, 3], [3, 3]]}

可选：Pandas简化实现 #

如果允许引入pandas依赖，代码可以更简洁：

import pandas as pd

coords = np.argwhere(labeled_array > 0)
df = pd.DataFrame(coords, columns=['y', 'x'])
df['label'] = labeled_array[coords[:,0], coords[:,1]]
result = df.groupby('label').agg(list).apply(tuple, axis=1).to_dict()

效率说明 #

• 逐标签调用np.where的时间复杂度为O(M*N)，M为标签数量，N为数组元素总数，标签越多性能下降越明显
• 上述向量化方案时间复杂度为O(N log N)，仅和数组总大小相关，和标签数量无关，标签数超过10个时性能就会显著超过循环方案，数组尺寸越大、标签越多优势越明显。
• 最后生成字典的推导式仅对标签级别的少量数据操作，开销可以忽略，不属于影响性能的循环逻辑。

内容的提问来源于stack exchange，提问作者Canol Gökel